This is a decidedly difficult pattern. Would it be possible to move taking the gradient (but not the optimizer step) to where the output is generated? That would make it much easier because you only need to keep the gradients around, not the computational graph. The problem here is that you can only take gradients between the forward and changing the weights (so 7 would fail because 5 changed the weights from what they were in 3). If this isn’t an option you could probably replace all the modules to call .clone() on all parameters before using them. This is a bit tedious because you need to re-implement/subclass all the nn.Modules you are using, but it would be feasible.
So the other thing we touched upon is that for this kind of pattern, you probably want to use a list to store the input and output tensors, not a single large tensor you index into.
Best regards
Thomas