Skip to content

Commit 4126a52

Browse files
update knn declaration (#236)
1 parent e1d99de commit 4126a52

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

csrc/cluster.h

+4-2
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@ CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
2020
torch::optional<torch::Tensor> optional_start,
2121
torch::optional<torch::Tensor> optional_end);
2222

23-
CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
24-
torch::Tensor ptr_y, int64_t k, bool cosine);
23+
CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
24+
torch::optional<torch::Tensor> ptr_x,
25+
torch::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
26+
int64_t num_workers);
2527

2628
CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
2729
torch::Tensor ptr_y);

0 commit comments

Comments
 (0)