Arjun_Gupta:

After many steps of loss_contrastive being at around ~10^8 and many gradient overflows (here the original loss is ~50), both losses become NaNs.


This doesn’t seem right. Are you expecting the loss to have such a high value?
If the loss gets NaN values, could you check, if your model is also outputting NaNs?
The zero division error is raised, as the loss scaler detects NaN gradients (due to an overflow, which might sometimes happen and is expected, or due to the NaN loss, which should never happen), and thus reduces the scaling factor until it’s not representable anymore.
That being said, we recommend to use the native amp implementation 17 instead of apex.