victorc25:

If I only autocast parts of the loss, wouldn’t calculations become inacurate?


That shouldn’t be the case, since autocast should take care of using FP16, where it’s safe to do so, and fallback to FP32 if necessary.
Could you post the setup for the loss function, so that I could run it in isolation?



 victorc25:

In the case of SSIM, the scale also drops heavily, but stabilizes at a scale of 2.


The initial loss scale might drop and stabilize, which is the expected behavior.
However, the outputs (and thus loss) should never contain any NaN values.