I see a few things:

If you want a function to be differentiable with autograd, you have to use only torch Tensor and torch ops. So if you extract things as python number or numpy array, it won’t work.

torch.tensor() creates a Tensor based on the raw data that is given to it. In particular, it will also not propagate gradient as it considers these data as plain numbers and not Tensors.