Let’s assume we want to do this in 0.4.

Do you really need a Function - if so, what would be your backward?
I would probably first switch to new-style (0.2 or so) Functions first. There is an example in the torch.autograd.Function 47 documentation. (Note that there is a subtle difference around variables vs. tensors if you look at 0.2 and 0.3.x docs, but let’s focus on 0.4).
I’m not quite sure that y.backward(y) - i.e. differentiating y w.r.t. the inputs while assuming that dloss/dy is y is something used very frequently.

Best regards
Thomas