[ROCm] fix nightly 5.6 build (pytorch#116029)

ROCm 5.6 nightly wheel build broken by pytorch#114329. This fixes it. Pull Request resolved: pytorch#116029 Approved by: https://github.com/huydhn, https://github.com/jithunnair-amd, https://github.com/atalman
dmenig · Dec 21, 2023 · 8e59208 · 8e59208
1 parent c78631c
commit 8e59208
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/aten/src/ATen/cuda/CUDABlas.cpp b/aten/src/ATen/cuda/CUDABlas.cpp
@@ -1102,7 +1102,7 @@ void int8_gemm(
 #endif // (!defined(USE_ROCM) && !defined(_MSC_VER)) || (defined(USE_ROCM) && ROCM_VERSION >= 50700)
 
 // ROCm 5.6 hipblas matches the const Dtype *A API, but prior hipblas does not.
-#if defined(USE_ROCM) && ROCM_VERSION <= 50600
+#if defined(USE_ROCM) && ROCM_VERSION < 50600
 #define ROCM_CONST_BUG
 #else
 #define ROCM_CONST_BUG const

diff --git a/aten/src/ATen/cuda/CublasHandlePool.cpp b/aten/src/ATen/cuda/CublasHandlePool.cpp
@@ -28,7 +28,7 @@ namespace at::cuda {
 
 namespace {
 
-#ifdef USE_ROCM
+#if defined(USE_ROCM) && ROCM_VERSION >= 50700
 void createCublasLtHandle(cublasLtHandle_t *handle) {
   TORCH_CUDABLAS_CHECK(cublasLtCreate(handle));
 }