Hi Augustijn,
the idea sounds really interesting and it would be great, if you could keep us updated about your progress here. 



 AuguB:

I understand that a call to with torch.nograd() i can do forward passes without autograd doing calculations in the background, so I had the idea to surround the second and third forward pass in the nograd() , but is this the best way to go here?


I think that’s the best approach to make sure only the output of the first forward pass has a valid grad_fn to calculate the gradients.
Let us know, if you encounter any problems with this approach.