Address review comments

Author: nicolas-chaulet

cpu/src/interpolate.cpp

@@ -30,7 +30,10 @@ at::Tensor knn_interpolate(at::Tensor features, at::Tensor idx, at::Tensor weigh
                 {
                     output_a[b][c][p] = 0;
                     for (int i = 0; i < idx.size(2); i++)
-                        output_a[b][c][p] += features_a[b][c][idx_a[b][p][i]] * weight_a[b][p][i];
+                    {
+                        auto new_idx = idx_a[b][p][i];
+                        output_a[b][c][p] += features_a[b][c][new_idx] * weight_a[b][p][i];
+                    }
                 }
             }
         }

cpu/src/neighbors.cpp

@@ -321,6 +321,6 @@ void nanoflann_knn_neighbors(vector<scalar_t>& queries, vector<scalar_t>& suppor
             neighbors_indices[i + current_pos] = ret_index[i];
             dists[i + current_pos] = out_dist_sqr[i];
         }
-        current_pos += nMatches;
+        current_pos += k;
     }
 }

torch_points/knn.py

@@ -1,5 +1,6 @@
 import torch_points.points_cpu as tpcpu
 
+
 def knn(pos_support, pos, k):
     """ Dense knn serach
     Arguments:
@@ -12,5 +13,6 @@ def knn(pos_support, pos, k):
         dist2 - [B,M,k] squared distances
     """
     assert pos_support.dim() == 3 and pos.dim() == 3
-    
-    return tpcpu.dense_knn(pos_support, pos, k)
+    if pos_support.is_cuda:
+        raise ValueError("CUDA version not implemented, use pytorch geometric")
+    return tpcpu.dense_knn(pos_support, pos, k)