fixed a bug in which max_num_neighbors was not handled correctly

rusty1s · rusty1s · commit 0841c7c30cc1 · 2020-08-05T07:05:42.000+02:00
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);
           }

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,8 @@ torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,`
`88`	`88`	`size_t num_matches = mat_index.index->radiusSearch(`
`89`	`89`	`y_data + i * y.size(1), r * r + 0.00001, ret_matches, params);`
`90`	`90`
`91`		`- for (size_t j = 0; j < num_matches; j++) {`
	`91`	`+ for (size_t j = 0;`
	`92`	`+ j < std::min(num_matches, (size_t)max_num_neighbors); j++) {`
`92`	`93`	`out_vec.push_back(x_start + ret_matches[j].first);`
`93`	`94`	`out_vec.push_back(i);`
`94`	`95`	`}`