Update deprecated CUB types (#851)

NVIDIA · Jan 30, 2025 · 6325970 · 6325970
1 parent 5010a55
commit 6325970
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/include/matx/transforms/cub.h b/include/matx/transforms/cub.h
@@ -993,7 +993,7 @@ inline void ExecSort(OutputTensor &a_out,
       if (a.IsContiguous()) {
         cub::DeviceSelect::If(d_temp,
                               temp_storage_bytes,
-                              cub::CountingInputIterator<index_t>{0},
+                              thrust::counting_iterator<index_t>(0),
                               a_out.Data(),
                               cparams_.num_found.Data(),
                               static_cast<int>(TotalSize(a)),
@@ -1004,7 +1004,7 @@ inline void ExecSort(OutputTensor &a_out,
         tensor_impl_t<typename InputOperator::value_type, InputOperator::Rank(), typename InputOperator::desc_type> base = a;
         cub::DeviceSelect::If(d_temp,
                               temp_storage_bytes,
-                              cub::CountingInputIterator<index_t>{0},
+                              thrust::counting_iterator<index_t>(0),
                               a_out.Data(),
                               cparams_.num_found.Data(),
                               static_cast<int>(TotalSize(a)),
@@ -1016,7 +1016,7 @@ inline void ExecSort(OutputTensor &a_out,
     else {
       cub::DeviceSelect::If(d_temp,
                             temp_storage_bytes,
-                            cub::CountingInputIterator<index_t>{0},
+                            thrust::counting_iterator<index_t>(0),
                             a_out.Data(),
                             cparams_.num_found.Data(),
                             static_cast<int>(TotalSize(a)),