Hi, welcome !
In general, to avoid this issue, we make (most) nn.Module stateless.
I think the problem you’re seeing is that if you run this for loop for a while, it will eat all your memory.
This is because, every new x you compute is given by x = transform(x, p).
At the very first iteration, x does not require gradient and p does.
So the backward sets the gradients of p and you’re happy.
But at the next iteration, the x that was the result of the op above does require gradients. And so at the second iteration, when you call backward, you will backward through both the second and the first iterations.
To prevent this, you want to break the link (in the autograd sense) between the x from one iteration to the other. You can do so with self.x = self.x.detach().
If you want to make sure that the graph of what is considered in the backward is what you expect, you can use a tool like torchviz 1 to plot this graph.
Hope this helps !