Hi,
If you use a custom Function, the backward should compute the gradient wrt the inputs by using the gradient wrt the outputs that are given to you.
As you don’t use any advanced option of np.gradient, it might just be simpler to re-implement it by hand using pytorch’s ops as you simply compare the adjacent values in the given input.