Another approach is to use torchviz: https://github.com/szagoruyko/pytorchviz/ 1
This can be used to plot the graph used to compute gradients.
You can use that just before the .backward() call.
You want to make sure that the graph printed by the first iteration and the second one are the same.
If they are not, then you want to find where the second one links to part of the first one to know where to add the .detach().
If Let me know if you need help understanding these graphs.