From 2b0aad181daff54cfbe6bd728dd22c4fd3398676 Mon Sep 17 00:00:00 2001
From: Seunghwa Kang <seunghwak@nvidia.com>
Date: Tue, 5 Dec 2023 17:59:45 -0800
Subject: [PATCH] leiden bug fix

---
 cpp/src/community/leiden_impl.cuh | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cpp/src/community/leiden_impl.cuh b/cpp/src/community/leiden_impl.cuh
index b6e20272de9..1e2b8f2ad44 100644
--- a/cpp/src/community/leiden_impl.cuh
+++ b/cpp/src/community/leiden_impl.cuh
@@ -568,17 +568,17 @@ void flatten_leiden_dendrogram(raft::handle_t const& handle,
   leiden_partition_at_level<vertex_t, multi_gpu>(
     handle, dendrogram, clustering, dendrogram.num_levels());
 
-  rmm::device_uvector<vertex_t> unique_cluster_ids(graph_view.number_of_vertices(),
+  rmm::device_uvector<vertex_t> unique_cluster_ids(graph_view.local_vertex_partition_range_size(),
                                                    handle.get_stream());
   thrust::copy(handle.get_thrust_policy(),
                clustering,
-               clustering + graph_view.number_of_vertices(),
+               clustering + graph_view.local_vertex_partition_range_size(),
                unique_cluster_ids.begin());
 
   remove_duplicates<vertex_t, multi_gpu>(handle, unique_cluster_ids);
 
   relabel_cluster_ids<vertex_t, multi_gpu>(
-    handle, unique_cluster_ids, clustering, graph_view.number_of_vertices());
+    handle, unique_cluster_ids, clustering, graph_view.local_vertex_partition_range_size());
 }
 
 }  // namespace detail