bug fix

rapidsai · Sep 19, 2023 · 52e2f57 · 52e2f57
1 parent ed7b1a4
commit 52e2f57
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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,