What are you doing with total_loss?
Currently you are storing the computation graph in it.
If you just need it for printing, you should use:
total_loss += loss.item()

Or do you need it somewhere for a backward pass?