From b86b38045dc8407158e166d2309c462aa7e8c8d5 Mon Sep 17 00:00:00 2001
From: hjwei <hjwei_xd@163.com>
Date: Mon, 23 Dec 2024 23:19:47 -0800
Subject: [PATCH] [BUG FIX]:fix get torch.version.cuda error when cuda is None
 in rocm

---
 op_builder/builder.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/op_builder/builder.py b/op_builder/builder.py
index 461281d4a569..f0ca90880d33 100644
--- 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',