From 4126a52680203425922bba34c730f57b1eb3372b Mon Sep 17 00:00:00 2001 From: Illia Ovcharenko Date: Thu, 14 Nov 2024 19:13:11 +0200 Subject: [PATCH] update knn declaration (#236) --- csrc/cluster.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/csrc/cluster.h b/csrc/cluster.h index e5a3b85..b615af6 100644 --- 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 optional_start, torch::optional 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 ptr_x, + torch::optional 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);