Hi,
There is no magic trick to force the autograd to only keep part of the graph I’m afraid.
If you want to avoid recomputing this, you can call .backward(retain_graph=True) which might increase you peak memory a little bit. Also if you do that, I would recommend you either:

Move the content of the while loop into a different function. To make sure that at the end of the loop, all the temporary variable are deleted before starting the next operation
Add a bunch of del foo at the end of your loop for outputs, all_grad, and loss. (again to make sure they don’t stay alive until the next iteration).