Thanks for the update! That’s strange indeed, as the scaler.step(optimizer) will skip the parameter updates if gradients are invalid and the parameters themselves are stored in float32 so cannot easily overflow.
Could you add the parameter check during the training and check which iteration creates the first invalid parameters? Once you are detecting them, print the gradients and check if any of them also contained invalid values, as these parameters should have never been updated. Also, I assume that you’ve verified the inputs and made sure they are not containing any invalid values?