[ROCm] fix for break on nightly whls

- Fix the break caused by PR#72854 Please review CC @jeffdaily @jithunnair-amd Pull Request resolved: pytorch#73399
lcy-seso · Feb 25, 2022 · 881a574 · 881a574
1 parent fd99944
commit 881a574
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/aten/src/ATen/cuda/cub.cuh b/aten/src/ATen/cuda/cub.cuh
@@ -51,7 +51,7 @@
 #define ROCM_HIPCUB(x) x
 #endif
 
-#if !CUB_SUPPORTS_NV_BFLOAT16() || \
+#if (!defined(USE_ROCM) && !CUB_SUPPORTS_NV_BFLOAT16()) || \
      (defined(USE_ROCM) && ROCM_VERSION >= 40500)
 
 #if !defined(USE_ROCM)