While training the first step, I can see that the loss is decreasing for each epoch, however at each epoch I see that vae is nan at epoch 1: 