In addition, i noticed that the PyTorch told me the gradients became NaN several iterations (about 40 with batch size 32) before inputs became NaN. I figure that the abnormal loss causes the weights to become NaN. So the main reason may be amp not working well.