[BUG FIX]:fix get torch.version.cuda error when cuda is None in rocm

microsoft · Dec 24, 2024 · b86b380 · b86b380
1 parent eea5304
commit b86b380
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/op_builder/builder.py b/op_builder/builder.py
@@ -839,7 +839,9 @@ def cxx_args(self):
 
         CPU_ARCH = self.cpu_arch()
         SIMD_WIDTH = self.simd_width()
-        CUDA_ENABLE = self.is_cuda_enable()
+        CUDA_ENABLE = (
+            "-D__DISABLE_CUDA__" if self.is_rocm_pytorch() else self.is_cuda_enable()
+        )
         args += [
             CPU_ARCH,
             '-fopenmp',