RuntimeError using {batch,event}_shape=[] base/target distributions

[feynmanliang]

Using a batch_shape=[] and event_shape=[] base/target results in a shape mismatch RuntimeError

import torch
import torch.distributions as dist
import flowtorch
import flowtorch.bijectors as bijectors
# Lazily instantiated flow plus base and target distributions
flow = bijectors.AffineAutoregressive(
    flowtorch.params.DenseAutoregressive()
)

########################
## LOOK HERE ##
base_dist = dist.Normal(0, 1)
target_dist = dist.Normal(5, 1)
#########################


# Instantiate transformed distribution and parameters
new_dist, params = flow(base_dist)
# Training loop
opt = torch.optim.Adam(params.parameters(), lr=1e-3)
for idx in range(501):
    opt.zero_grad()
    # Minimize KL(p || q)
    y = target_dist.sample((1000,))
    loss = -new_dist.log_prob(y).mean()
    if idx % 100 == 0:
        print('epoch', idx, 'loss', loss)
        
    loss.backward()
    opt.step()

sns.relplot(
    data=pd.DataFrame(new_dist.sample((100,)).detach().numpy()),
    x=0, y=1
)