Fix include for <iterator> (#1929)

Follow-up fix for PR #1863
NVIDIA · Jul 2, 2024 · 7f2c1d1 · 7f2c1d1
1 parent 327420b
commit 7f2c1d1
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/cub/cub/util_type.cuh b/cub/cub/util_type.cuh
@@ -46,7 +46,6 @@
 #include <cub/detail/uninitialized_copy.cuh>
 
 #include <cuda/std/cstdint>
-#include <cuda/std/iterator>
 #include <cuda/std/limits>
 #include <cuda/std/type_traits>
 
@@ -59,6 +58,12 @@ _CCCL_DIAG_POP
 #  endif // !_CCCL_CUDACC_BELOW_11_8
 #endif // _CCCL_HAS_NV_BF16
 
+#ifdef _CCCL_COMPILER_NVRTC
+#  include <cuda/std/iterator>
+#else // !defined(_CCCL_COMPILER_NVRTC)
+#  include <iterator>
+#endif // defined(_CCCL_COMPILER_NVRTC)
+
 CUB_NAMESPACE_BEGIN
 
 #ifndef CUB_IS_INT128_ENABLED