Hi @YASJAY, I should’ve clarified:  I’m running this job on a GPU cluster so I’m confident memory is completely free prior to job submission. My batch size cannot be changed at all, but it is fairly small, four.
My issue is that I quickly run out of memory, likely because of the total += torch.exp(soft_inner_logits) step; I believe that keeps piling on computation graphs on top of each other, using up a lot of memory. I was more looking for ways to re-write what I’ve done in a more memory-efficient way: E.g.: backpropagating in the loop to free up intermediary memory (Which unfortunately I can’t do in this case since the loss cannot be calculated until I’ve gone through N passes).