How to calculate loss and gradient of index

I have a misaligned point cloud and rotation matrix and I want to use backprop to tune the rotation matrix in 2d plane.

The problem is in the forward all I can get from the rotation matrix is index in the 2d plane which is duplicated like.

>>> x
[500.99, 500.5, 411.2, 411.02, 411]

But the desired output for this case should be

[500, 411]

While the ground truth maybe

[505, 501, 460, 420]

I can’t compare the float x and ground truth using mse directly because the dimension is not equal.

Both the x and ground truth are index of a 512 dimension vector but the gradient disappear as soon as I use x.long() so I can’t use them as index either.

How to properly calculate loss of this problem?

Thanks