lib/knn/src/knn_pytorch.h

int knn(THCudaTensor *ref_tensor, THCudaTensor *query_tensor,

    THCudaLongTensor *idx_tensor);