bluesky314:

I get NaN loss from the first batch continuing my trained model.


This would point towards an overflow in some intermediate activations or are you also seeing a NaN output if you train the model from scratch?



 a_d:

Remove the scaler.scale(loss) and simply run loss.backward() outside the autocast, that might fix it…


No, the scaler is an essential part of mixed-precision training and should not be removed.
Here 70 you can see the typical training loop with comments for each step.