Thanks for the code.
If I set use_mixed_precision=False, the code also outputs NaNs, so it seems your training blows up:
iteration 0  nan found  False
Loss is  10132.611328125
iteration 1  nan found  True
Loss is  inf
iteration 2  nan found  True
Loss is  nan
iteration 3  nan found  True
Loss is  nan
iteration 4  nan found  True
Loss is  nan
iteration 5  nan found  True
Loss is  nan
iteration 6  nan found  True
Loss is  nan
iteration 7  nan found  True
Loss is  nan
iteration 8  nan found  True
Loss is  nan

I guess the learning rate is just too high (for AMP and vanilla FP32 training).
If you use reduction='mean' or lower the learning rate, the model seems to work fine in both modes (for at least 60 epochs, as I’ve stopped it afterwards).