Solved by ptrblck in post #3 
              
              
                You can define a custom autograd.Function as described here and use it in a custom nn.Module by falling it inside the forward method of the module.