How to compute a block of code using full precision when training with AMP?
I have encountered a few cases so far where some losses (such as the ms-ssim implementation from VainF on GitHub ) that used to work under full precision training will no longer work when using AMP.
I would like to know if there is any way to locally disable AMP for just a few steps of computation, either through Lightning or through native PyTorch. (It seems that torch.cuda.amp.autocast(enabled=False) might be able to do the job, but I am worried about it messing with the Lightning internals)
Since I calculate a variety of losses before combining them, all in the same step, I am thinking that doing so will still allow me to leverage some of the advantages from AMP