I’ve now committed the code with AMP to the branch (https://github.com/victorca25/BasicSR/tree/dev2/codes 3).
I’m giving up on making AMP work properly after that convolution in the CX loss, I can’t find the cause of the problem, so I’m forcing the tensors to be converted to FP32 right after that convolution as a workaround for now (https://github.com/victorca25/BasicSR/blob/57577667067c99194098ce034a97d0bef028620d/codes/models/modules/loss.py#L707 4).
It runs and the other losses work fine. The forward pass also works fine in the autocast context, so at least its better than nothing.