Use AT_ROCM_ENABLED() instead of USE_ROCM

IvanYashchuk · IvanYashchuk · commit 42b868c2bda6 · 2022-03-19T12:56:28.000-05:00
diff --git a/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp b/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp
@@ -3231,7 +3231,7 @@ void lstsq_kernel(const Tensor& a, Tensor& b, Tensor& /*rank*/, Tensor& /*singul
         "Please rebuild with cuSOLVER.");
 #endif
   } else { // m >= n
-#if !defined(USE_ROCM)
+#if !AT_ROCM_ENABLED()
     // On CUDA platform we use either cuBLAS or cuSOLVER here
     // the batched vs looped dispatch is implemented based on the following performance results
     // https://github.com/pytorch/pytorch/pull/54725#issuecomment-832234456
@@ -3244,7 +3244,7 @@ void lstsq_kernel(const Tensor& a, Tensor& b, Tensor& /*rank*/, Tensor& /*singul
     // On ROCm platform we can only use MAGMA here
     // If MAGMA is not available, an error will be thrown
     gels_magma(a, b, infos);
-#endif // USE_ROCM()
+#endif // !AT_ROCM_ENABLED()
   }
 }
 
