You could try to replace loss += with loss = loss +.



 mohela:

I also had to have loss.backward(retain_graph=True) in the batch loop.


Was that a surprise or expected?
The Thomas rule of thumb here is: In most cases, retain_graph=True is the wrong solution to the “attempting to backward for a second time” error. I always try to articulate why I would need to keep around intermediate results before believing I should.
Best regards
Thomas