Solved by alepack in post #2 
              
              
                Hi, 
I think that if you write a simple function like 
import torch 
def sqrt_act_fun(x): 
num = x 
den = torch.sqrt(1.0 + torch.mul(x,x)) 
return torch.div(num, dev) 
Provided that “x” is a torch tensor and that you use the output to compute (directly or indirectly) the loss value L, then the gradi…