pyro.clear_param_store()
optim = pyro.optim.Adam({"lr": 1e-3})
elbos = []
def callback(bnn, i, e):
    elbos.append(e)
    
with tyxe.poutine.local_reparameterization():
    bnn.fit(loader, optim, 10000, callback)

...
C:\Miniconda3\envs\pta\lib\site-packages\torch\distributions\transforms.py in _inv_call(self, y)
    155         """
    156         if self._cache_size == 0:
--> 157             return self._inverse(y)
    158         x_old, y_old = self._cached_x_y
    159         if y is y_old:

C:\Miniconda3\envs\pta\lib\site-packages\torch\distributions\transforms.py in _inverse(self, y)
    430 
    431     def _inverse(self, y):
--> 432         if y.dim() < self.codomain.event_dim:
    433             raise ValueError("Too few dimensions on input")
    434         return self.base_transform.inv(y)

AttributeError: 'PyroParam' object has no attribute 'dim'