Hi,
I am not as familiar with the jit c++ API.
But did you make sure that the input require gradients? And if it is not a leaf, you can call .retain_grad() on it to ensure its .grad field is properly populated.
Also make sure that you don’t disable the autograd with a NoGradGuard.