bug fix

rapidsai · Aug 20, 2024 · 9625e0c · 9625e0c
1 parent 8193a91
commit 9625e0c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cpp/src/prims/detail/per_v_transform_reduce_e.cuh b/cpp/src/prims/detail/per_v_transform_reduce_e.cuh
@@ -1160,7 +1160,8 @@ void per_v_transform_reduce_e_edge_partition(
         auto segment_key_last  = edge_partition_key_last;
         if constexpr (use_input_key) {
           segment_key_first += (*key_segment_offsets)[3];
-          segment_key_last += (*key_segment_offsets)[4];
+          segment_key_last =
+            segment_key_first + ((*key_segment_offsets)[4] - (*key_segment_offsets)[3]);
         } else {
           assert(segment_key_first == nullptr);
           assert(segment_key_last == nullptr);