loss and mse are nan in experiments heat-diffusion

[Lucidreamer9]

Hi,

I tried the heat-diffusion experiments, I generated data using [gaussian_random_fields.ipynb] and trained a model using python train.py --name HeatDiffusion --train-file data/files/32/2d_gaussian_random_field_4.h5 --test-file data/files/32/2d_gaussian_random_field_4_test.h5 --batch-size 4 EmbeddedPhysics

However, the loss becomes nan during training, and the output is: Test MSE probability flow: nan
Test MSE reverse sde: nan