(Fix an issue of wrong variable name)

rapidsai · Mar 26, 2024 · 440d97d · 440d97d
1 parent 7cd6f0e
commit 440d97d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh
@@ -90,12 +90,12 @@ void search_main_core(
   CagraSampleFilterT sample_filter = CagraSampleFilterT())
 {
   RAFT_LOG_DEBUG("# dataset size = %lu, dim = %lu\n",
-                 static_cast<size_t>(index.data().n_rows()),
-                 static_cast<size_t>(index.data().dim()));
+                 static_cast<size_t>(graph.data().n_rows()),
+                 static_cast<size_t>(graph.data().dim()));
   RAFT_LOG_DEBUG("# query size = %lu, dim = %lu\n",
                  static_cast<size_t>(queries.extent(0)),
                  static_cast<size_t>(queries.extent(1)));
-  RAFT_EXPECTS(queries.extent(1) == dataset_desc.dim, "Queries and index dim must match");
+  RAFT_EXPECTS(queries.extent(1) == dataset_desc.dim, "Queries and graph dim must match");
   const uint32_t topk = neighbors.extent(1);
 
   cudaDeviceProp deviceProp = resource::get_device_properties(res);

diff --git a/cpp/include/raft/neighbors/detail/cagra/graph_core.cuh b/cpp/include/raft/neighbors/detail/cagra/graph_core.cuh
@@ -696,6 +696,10 @@ void mst_optimization(raft::resources const& res,
       // If the number of clusters does not converge to 1, then edges are
       // made from all nodes not belonging to the main cluster to any node
       // in the main cluster.
+      raft::copy(
+        cluster_size_ptr, d_cluster_size_ptr, (size_t)graph_size, resource::get_cuda_stream(res));
+      raft::copy(label_ptr, d_label_ptr, (size_t)graph_size, resource::get_cuda_stream(res));
+      resource::sync_stream(res);
       uint32_t main_cluster_label = graph_size;
 #pragma omp parallel for
       for (uint64_t i = 0; i < graph_size; i++) {