The documentation has an example for a simple exponential function here
Another example from the forums is here, although it’s for a custom Double Backward (so 2nd derivative) it might help illustrate what the custom torch.autograd.Function
does and what you need code in for it to work.