I don’t think this would be possible, as



 Kevin_Xue:

The indices are the pixel location, which has requires_grad and is computed from the leaf tensor.


would be disallowed, wouldn’t it?
You could probably try to come up with a custom autograd.Function, which could define (somehow) the backward and gradients for indices although I wouldn’t have a good idea if and how this could work.