Hi,
Just as a note, check the doc 8 for register backward hook. In particular the warning that it might not return what you want in some cases at the moment.
The memory problem here is that you do total_grad_input.append(grad_input). So that list will contain more and more Tensors. You want to make sure to remove them when you don’t need them anymore.