Solved by albanD in post #2 
              
              
                Hi, 
You will need to use a custom autograd.Function for that. This section of the doc shows you how to do so and how to use it.