I think it’s that when you accumulate the loss, you should unpack it, like loss.data[0]
Otherwise the graph will never be freed.
See this discussion: CUDA memory continuously increases when net(images) called in every iteration