Hello again!
Thanks to your suggestion of autocasting only parts of the loss, I found the exact line causing the scale to drop, its this one: https://github.com/victorca25/BasicSR/blob/3fb5851af3a411df153f0b0e5873f3b2de324a3d/codes/models/modules/loss.py#L704 2
dist = F.conv2d(I_features_i, T_features_i).permute(0, 2, 3, 1).contiguous()
If just this line is put inside an autocast context, the scale is halved every iteration.
Any suggestions on what could be the reason? It doesn’t appear to be anything special there.