@ptrblck Sorry to bother, do you know if there are any ways to capture memory leaks when using grad() from autograd.?
For instance in the MWE in my first post if I replace mse(x @ dummy_psi * dummy_w) with mse(vgg_like(x) * dummy_w) I get cuda error out of memory.
When I switch the whole example to cpu I can monitor the ram keep increasing on every iteration until it takes all my available ram. I’ve tracked down the issue to the grad() method because when I comment it out everything works fine but I don’t know why this is happening or where the leak commes from?
Also this leak happens regardless of using retain-graph=True or create_graph=True.