How to try/catch errors during training

Hi,

I think it is a “known” issue with python exceptions. See Exception leaks in Python 2 and 3 | Kristján's Cosmic Percolator
In your case, since the differentiable output is in the current frame, it is kept alive by the exception as so holds on to the GPU memory forever since you never exit the function.

I think a simple fix for you would be to move the content of the loop in a separate function that will be exited at every iteration (clearing the exception properly).