Hi,
You can check the doc on creating custom Functions for more details: https://pytorch.org/docs/stable/notes/extending.html 18
But the gist is that if you need to save input/output of the forward, you have to use ctx.save_for_backward(Values) and get them back as Values, = ctx.saved_tensors otherwise you will see these leaks.