MNIST Addition example

[thiviyanT]

Hello @EleMisi !

I encountered the following error when attempting to run VAEL on the MNIST addition example, so I thought it would be best to report it here.

Traceback (most recent call last):
  File "/content/VAEL/run_VAEL.py", line 37, in <module>
    run_mnist_vael(config_object['exp_config'], exp_class, exp_folder, data_folder, n_digits=config_object['n_digits'],
  File "/content/VAEL/utils/mnist_utils/mnist_task_VAEL.py", line 313, in run_mnist_vael
    checkpoint_path, epoch, train_info, validation_info = train_PLVAE(model,
  File "/content/VAEL/utils/mnist_utils/train.py", line 124, in train_PLVAE
    loss, recon_loss, gauss_kl_div, queryBCE, labelBCE = loss_function(recon_batch,
  File "/content/VAEL/utils/mnist_utils/train.py", line 19, in loss_function
    recon_loss = - img_log_likelihood(x_recon, x).mean()
  File "/content/VAEL/utils/mnist_utils/train.py", line 55, in img_log_likelihood
    return torch.distributions.Laplace(recon, torch.ones_like(recon)).log_prob(xs).sum(dim=(1, 2, 3))
  File "/usr/local/lib/python3.10/dist-packages/torch/distributions/laplace.py", line 52, in __init__
    super().__init__(batch_shape, validate_args=validate_args)
  File "/usr/local/lib/python3.10/dist-packages/torch/distributions/distribution.py", line 68, in __init__
    raise ValueError(
ValueError: Expected parameter loc (Tensor of shape (30, 1, 28, 56)) of distribution Laplace(loc: torch.Size([30, 1, 28, 56]), scale: torch.Size([30, 1, 28, 56])) to satisfy the constraint Real(), but found invalid values:
tensor([[[[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]]],


        [[[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]]],


        [[[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]]],


        ...,


        [[[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]]],


        [[[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]]],


        [[[nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          ...,
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan],
          [nan, nan, nan,  ..., nan, nan, nan]]]], device='cuda:0',
       grad_fn=<SigmoidBackward0>)
 49% 692/1400 [00:06<00:06, 103.14it/s]

I followed the instruction on the README. I am it running on a Linux machine. Using pytorch 2.1.0 with cuda 12.1.