fixed a bug in which max_num_neighbors was not handled correctly

rusty1s · Aug 5, 2020 · 0841c7c · 0841c7c
1 parent 99d0792
commit 0841c7c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/csrc/cpu/radius_cpu.cpp b/csrc/cpu/radius_cpu.cpp
@@ -88,7 +88,8 @@ torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
           size_t num_matches = mat_index.index->radiusSearch(
               y_data + i * y.size(1), r * r + 0.00001, ret_matches, params);
 
-          for (size_t j = 0; j < num_matches; j++) {
+          for (size_t j = 0;
+               j < std::min(num_matches, (size_t)max_num_neighbors); j++) {
             out_vec.push_back(x_start + ret_matches[j].first);
             out_vec.push_back(i);
           }