Loss value appears NAN

[qiqiing]

Hello, when I use the cross-scale attention module in your code and my own data for training, the Loss value becomes NAN as the number of training iterations increases. I can make sure that there is no NAN data in my data. When I reduce the size of the training image or set the batchsize to 1, the training loss becomes normal. Is it because the cross-scale attention module may cause gradient explosion? How can I improve it?

[HarukiYqM]

Hi, I am not quite sure why this happen but I do not observe such effects when training on DIV2K. It okay to reduce the input size, e.g., from 48 to 32 while keep the model working well. But one simple strategy is to resume from the checkpoint before it becomes NaN, or reduce the learning rate, for example, starting from 5e-5.
Hope this can solve your problem.