diff --git a/cpp/include/raft/cluster/detail/kmeans_balanced.cuh b/cpp/include/raft/cluster/detail/kmeans_balanced.cuh
index da40ef4e0c..3b5d3ff02a 100644
--- 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