NAN loss during training

[hao-pt]

Thank for your great work! However, I got NAN loss issue when training with your provided code! As I saw in the paper, you also have an additional GAN loss which is not supported in the current repo. Does this lead to NAN loss when this type of loss is missing? I use the same training config and train on DIV2K.