Hi,
You can find on this page on the doc https://pytorch.org/docs/stable/notes/extending.html 12 all the information on how to define a custom elementary function for the autograd for which you define both the forward and the backward.