Any updates on whether this is possible? The context is the following: I want to implement a custom backward function for an invertible layer. Instead of saving the activations during the forward pass, the backward function takes the original output y as and additional input argument, from which the input x is recomputed and used to compute gradients. I therefore need my custom backward function to not only accept the gradient as an input, but also the original output (in the example of @mathkobe this is the additional argument P).
I added the additional argument to the backward functions in torch/tensor.py and torch/autograd/__init__.py. But after that, the functions step into the C code and I cannot follow them with the Python debugger.
Any help or pointers on how to extend this would be great! Thanks a lot!