diff --git a/cpu/include/neighbors.h b/cpu/include/neighbors.h index 1c83851..32de50e 100644 --- a/cpu/include/neighbors.h +++ b/cpu/include/neighbors.h @@ -5,19 +5,18 @@ #include #include -using namespace std; template -int nanoflann_neighbors(vector& queries, vector& supports, - vector& neighbors_indices, vector& dists, float radius, +int nanoflann_neighbors(std::vector& queries, std::vector& supports, + std::vector& neighbors_indices, std::vector& dists, float radius, int max_num, int mode, bool sorted); template -int batch_nanoflann_neighbors(vector& queries, vector& supports, - vector& q_batches, vector& s_batches, - vector& neighbors_indices, vector& dists, +int batch_nanoflann_neighbors(std::vector& queries, std::vector& supports, + std::vector& q_batches, std::vector& s_batches, + std::vector& neighbors_indices, std::vector& dists, float radius, int max_num, int mode, bool sorted); template -void nanoflann_knn_neighbors(vector& queries, vector& supports, - vector& neighbors_indices, vector& dists, int k); +void nanoflann_knn_neighbors(std::vector& queries, std::vector& supports, + std::vector& neighbors_indices, std::vector& dists, int k); diff --git a/cpu/src/knn.cpp b/cpu/src/knn.cpp index 199e931..516c4c0 100644 --- a/cpu/src/knn.cpp +++ b/cpu/src/knn.cpp @@ -45,8 +45,8 @@ std::pair dense_knn(at::Tensor support, at::Tensor query CHECK_CPU(support); int b = query.size(0); - vector batch_idx; - vector batch_dist; + std::vector batch_idx; + std::vector batch_dist; for (int i = 0; i < b; i++) { auto out_pair = _single_batch_knn(support[i], query[i], k); diff --git a/cpu/src/neighbors.cpp b/cpu/src/neighbors.cpp index 3a7556c..7a58945 100644 --- a/cpu/src/neighbors.cpp +++ b/cpu/src/neighbors.cpp @@ -5,6 +5,10 @@ #include #include +using std::vector; +using std::pair; +using std::max; + template int nanoflann_neighbors(vector& queries, vector& supports, vector& neighbors_indices, vector& dists, float radius,