When training readout, loss is Nan

[lipanUestc]

First of all, thank you very much for your open source code! Your work is great!
But when I ran the code in "pytorch_IBA_train_readout.ipynb", I found information_ Loss is all Nan after the first epoch. Is there something wrong with me? Could you give me some advice?

[lipanUestc]

I found that the reason for this problem was that I changed the dataset. Is this problem caused by the estimator's evaluation of mean and variance? What should I do when I want to change the dataset?