Where do you print d.grad?
If I just add a print statement after adv_distance.backward(), I’ll get a valid gradient for d.
Also, if I register d as an attribute, I can print the gradient successfully.