soumickmj:

if I put the forward pass of only that module with autocast False, then it works for 14 iterations and then throughs the exact same error.


This points towards an instability on your overall training and doesn’t seem to be related to AMP.
I’m not sure how “Its only showing up when I’m using AMP.” fits into the new post.
For further debugging I would suggest to disable amp for now and check the inputs to the criterion, which create the invalid outputs/gradients.