update knn declaration (#236)

rusty1s · Nov 14, 2024 · 4126a52 · 4126a52
1 parent e1d99de
commit 4126a52
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/csrc/cluster.h b/csrc/cluster.h
@@ -20,8 +20,10 @@ CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
                    torch::optional<torch::Tensor> optional_start,
                    torch::optional<torch::Tensor> optional_end);
 
-CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
-                  torch::Tensor ptr_y, int64_t k, bool cosine);
+CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
+                  torch::optional<torch::Tensor> ptr_x,
+                  torch::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
+                  int64_t num_workers);
 
 CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
                       torch::Tensor ptr_y);