Set this flag for CUDAOpBuilder too

microsoft · Nov 7, 2023 · f65a857 · f65a857
1 parent d4e9bd9
commit f65a857
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/op_builder/builder.py b/op_builder/builder.py
@@ -597,12 +597,7 @@ def is_compatible(self, verbose=True):
         return super().is_compatible(verbose)
 
     def builder(self):
-        try:
-            if not self.is_rocm_pytorch():
-                assert_no_cuda_mismatch(self.name)
-            self.build_for_cpu = False
-        except BaseException:
-            self.build_for_cpu = True
+        self.build_for_cpu = not torch.cuda.is_available()
 
         if self.build_for_cpu:
             from torch.utils.cpp_extension import CppExtension as ExtensionBuilder