Using Mixed Precision but calculate loss using Full Precision

[JoakimHaurum]

Hi,

I am currently training a network for multi label classification.
For this I am using the Asymmetric Loss from Alibaba (https://github.com/Alibaba-MIIL/ASL)
However, when I train with mixed precision the loss goes to NaN.

This is a known problem (Alibaba-MIIL/ASL#53) and has been circumvented by calculating the loss in full precision 32-bit, while during the rest of the forward pass in mixed precision.

Is this possible with pytorch lightning or should I refactor the loss implementation?

[tchaton]

Dear @JoakimHaurum,

Yes, you can use torch.cuda.autocast(enabled=False) around the forward of your loss computation.

Also, we just merge support for bfloat16 which is more stable than float16.

It can be used with Trainer(precision="b16") on master + Torch nightly.

Best,
T.C