Weights precisions do not match up. AMP
Hi, thanks for your work!
My situation: I have a loss function that is defined in a separate file and it has weights and parameters and I import it into my pl module like
from arc_margin_loss import arcface_loss
self.loss = arcface_loss(parametes)

So when I try to use amp (Pytorch 1.6) I get those errors that I try to compare float and half tensors.
My question is how do I avoid this problem correctly?
Should I just add @autocast() decorator to the loss forward?
Or maybe I should I define loss inside pl module?
Thanks in advance!