I am using mixed precision training with a scaler and autocast(), my main reason to use fp16 is so I can fit a bigger model onto my gpu.
However it seems like the mixed precision training requires more GPU memory than without, as my code is now running out of GPU memory on the backward call where it could train before…
scaler.scale(loss).backward()

I am assuming that this is because of the copying of the fp32 weights when converting to fp16 in the backward call.
I followed https://pytorch.org/docs/stable/amp.html 5 and the other resources available but haven’t been able to find a solution to this, ideally I would like to have the copied fp32 weights on the cpu, especially if I am not using them at that point.
If I lower my (already small) batchsize from 4 to 2, it does work, however I wanted to use mixed precision training to increase my batch size/model size!
Any help would be much appreciated!