Solved by albanD in post #7 
              
              
                I don’t think you need a custom backward here no. You should simply update your function to only work with Tensor all the way and do not “repack” them with torch.tensor() or call .detach() on them as these would break the autograd link.