diff --git a/cpp/src/sampling/sampling_post_processing_impl.cuh b/cpp/src/sampling/sampling_post_processing_impl.cuh
index ff8da72ff35..0c38048f1fb 100644
--- a/cpp/src/sampling/sampling_post_processing_impl.cuh
+++ b/cpp/src/sampling/sampling_post_processing_impl.cuh
@@ -1743,10 +1743,11 @@ sort_sampled_edgelist(
                  (*edgelist_label_hop_offsets).begin(),
                  (*edgelist_label_hop_offsets).end(),
                  size_t{0});
-    thrust::for_each(
+    thrust::transform(
       handle.get_thrust_policy(),
       thrust::make_counting_iterator(size_t{0}),
       thrust::make_counting_iterator(num_labels * num_hops),
+      (*edgelist_label_hop_offsets).begin(),
       [edgelist_label_offsets = edgelist_label_offsets
                                   ? thrust::make_optional(std::get<0>(*edgelist_label_offsets))
                                   : thrust::nullopt,