Tensor size mismatch in VariationalMeanField.forward

[ischurov]

I'm trying to train model with variational: mean-field and get the following error:

user@maccie:prj/variational-autoencoder ‹master*›$ python3 train_variational_autoencoder_pytorch.py --variational mean-field
train_variational_autoencoder_pytorch.py:184: YAMLLoadWarning: calling yaml.load() without Loader=... is deprecated, as the default Loader is unsafe. Please read https://msg.pyyaml.org/load for full details.
  dictionary = yaml.load(config)
step:   0       train elbo: -557.94
Traceback (most recent call last):
  File "train_variational_autoencoder_pytorch.py", line 233, in <module>
    valid_elbo, valid_log_p_x = evaluate(cfg.n_samples, model, variational, valid_data)
  File "train_variational_autoencoder_pytorch.py", line 173, in evaluate
    elbo = log_p_x_and_z - log_q_z
RuntimeError: The size of tensor a (128) must match the size of tensor b (512) at non-singleton dimension 1