Hi,
Thank you for your reply! Here is my test code.
import torch
from nnd import NNDModule
data1 = torch.rand((2, 3, 5), requires_grad=True).cuda() #[b, c, n]
data2 = torch.rand((2, 3, 4), requires_grad=True).cuda() #[b, c, m]
dist = NNDModule()
print(torch.autograd.gradcheck(dist.double(), (data1.double(), data2.double()), eps=1e-6))

My function can only receive float precision input, so i use .double() to make it can receive double precision tensors, which i don’t know if that is a right way but the code has not raise the input error at least.Then, i input double precision random tensors to test, and the jacobin error raise .
The formula as the backward pass:
\[\begin{array}{l}
d{x_i} = 2*gradinput*({x_i} - {y_{xi}})\\
d{y_{xi}} =  - 2*gradinput*({x_i} - {y_{xi}})
\end{array}\]

the correspond code:
float g = grad_dist1[i]*2;
int id = idx1[i]; // the nearest neighbor index of xyz1[i] in xyz2
			
for (int k=0;k<c;k++) // traverse the channels
{
     // n is the number of item in xyz1
    // m is the number of item in xyz2
	atomicAdd(grad_xyz1 + i + k*n, g*(xyz1[i + k*n]-xyz2[id + k*m]));
	atomicAdd(grad_xyz2 + id + k*m, -(g*(xyz1[i + k*n]-xyz2[id + k*m])));
}

Also, i think the nearest neighbor search is not continuously differentiable because i think the max is not continuously differentiable，but  the torch.max() can also pass the gradcheck , so i’m a little confused