Solved by rustagiadi95 in post #2 
              
              
                Hello There, 
There is an equivalent of this functionality in PyTorch. 
https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html 
This article shows how to define a custom function with autograd tracking. You have to define the forward and backward pass of the func…