[CUDA] fixing partition bug (dmlc#6001)

DominikaJedynak · Mar 12, 2024 · c7f47d5 · c7f47d5
1 parent 9a2baff
commit c7f47d5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/partition/cuda/partition_op.cu b/src/partition/cuda/partition_op.cu
@@ -329,14 +329,14 @@ std::pair<IdArray, NDArray> GeneratePermutationFromRemainder(
     CUDA_CALL(cub::DeviceHistogram::HistogramEven(
         nullptr, hist_workspace_size, proc_id_out.get(),
         reinterpret_cast<AtomicCount*>(out_counts), num_parts + 1,
-        static_cast<IdType>(0), static_cast<IdType>(num_parts + 1),
+        static_cast<IdType>(0), static_cast<IdType>(num_parts),
         static_cast<int>(num_in), stream));
 
     Workspace<void> hist_workspace(device, ctx, hist_workspace_size);
     CUDA_CALL(cub::DeviceHistogram::HistogramEven(
         hist_workspace.get(), hist_workspace_size, proc_id_out.get(),
         reinterpret_cast<AtomicCount*>(out_counts), num_parts + 1,
-        static_cast<IdType>(0), static_cast<IdType>(num_parts + 1),
+        static_cast<IdType>(0), static_cast<IdType>(num_parts),
         static_cast<int>(num_in), stream));
   }
 
@@ -504,14 +504,14 @@ std::pair<IdArray, NDArray> GeneratePermutationFromRange(
     CUDA_CALL(cub::DeviceHistogram::HistogramEven(
         nullptr, hist_workspace_size, proc_id_out.get(),
         reinterpret_cast<AtomicCount*>(out_counts), num_parts + 1,
-        static_cast<IdType>(0), static_cast<IdType>(num_parts + 1),
+        static_cast<IdType>(0), static_cast<IdType>(num_parts),
         static_cast<int>(num_in), stream));
 
     Workspace<void> hist_workspace(device, ctx, hist_workspace_size);
     CUDA_CALL(cub::DeviceHistogram::HistogramEven(
         hist_workspace.get(), hist_workspace_size, proc_id_out.get(),
         reinterpret_cast<AtomicCount*>(out_counts), num_parts + 1,
-        static_cast<IdType>(0), static_cast<IdType>(num_parts + 1),
+        static_cast<IdType>(0), static_cast<IdType>(num_parts),
         static_cast<int>(num_in), stream));
   }