From 87efa982e907df1e9f3ee61308e9b651dc10be21 Mon Sep 17 00:00:00 2001
From: Elias Stehle <3958403+elstehle@users.noreply.github.com>
Date: Sat, 11 Jan 2025 09:27:28 -0800
Subject: [PATCH] testing seg type

---
 cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh
index 23ea5d086e..ebe3cb4aef 100644
--- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh
+++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh
@@ -76,7 +76,7 @@ namespace detail
 namespace segmented_sort
 {
 // Type used to index within segments within a single invocation
-using per_invocation_segment_offset_t = unsigned int;
+using per_invocation_segment_offset_t =  int;
 // Type used for total number of segments and to index within segments globally
 using global_segment_offset_t = ::cuda::std::int64_t;
 
@@ -576,7 +576,7 @@ template <typename LargeSegmentPolicyT,
 CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortContinuation(
   LargeKernelT large_kernel,
   SmallKernelT small_kernel,
-  int num_segments,
+  unsigned int num_segments,
   KeyT* d_current_keys,
   KeyT* d_final_keys,
   detail::device_double_buffer<KeyT> d_keys_double_buffer,