I don’t think the computations in the training code should affect the backward pass because I made sure to move everything onto the CPU. By the way, I also tried using torch.cuda.empty_cache() to clear the cache after every batch. It still didn’t help.
The rest of the tensor sizes are fixed to known dimensions. The only thing that is standing out to me as weird is that the values of the losses themselves are very very large. Is there any way that this is what is causing the issue somehow? Do you have any suggestions about what I could possibly do to fix this or find a work around to avoid this issue?