Another one

rapidsai · Dec 8, 2023 · 1075df5 · 1075df5
1 parent e9b5e00
commit 1075df5
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/cpp/include/raft/cluster/detail/kmeans_balanced.cuh b/cpp/include/raft/cluster/detail/kmeans_balanced.cuh
@@ -1016,7 +1016,7 @@ void build_hierarchical(const raft::resources& handle,
                    mesocluster_labels_buf.data(),
                    mesocluster_sizes_buf.data(),
                    mapping_op,
-                   &managed_memory,
+                   device_memory,
                    dataset_norm);
   }
 
@@ -1059,11 +1059,11 @@ void build_hierarchical(const raft::resources& handle,
                                              fine_clusters_nums_max,
                                              cluster_centers,
                                              mapping_op,
-                                             device_memory,
+                                             &managed_memory,
                                              device_memory);
   RAFT_EXPECTS(n_clusters_done == n_clusters, "Didn't process all clusters.");
 
-  rmm::device_uvector<CounterT> cluster_sizes(n_clusters, stream, &managed_memory);
+  rmm::device_uvector<CounterT> cluster_sizes(n_clusters, stream, device_memory);
   rmm::device_uvector<LabelT> labels(n_rows, stream, device_memory);
 
   // Fine-tuning k-means for all clusters