Avoid using cuda when it is installed but torch was compiled against …

…ROCM/hip (#1164)
facebookresearch · Dec 17, 2024 · f0a401c · f0a401c
1 parent 1f2b5a7
commit f0a401c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -438,7 +438,7 @@ def get_extensions():
     use_pt_flash = False
 
     if (
-        (torch.cuda.is_available() and ((CUDA_HOME is not None)))
+        (torch.cuda.is_available() and (CUDA_HOME is not None) and (torch.version.cuda is not None))
         or os.getenv("FORCE_CUDA", "0") == "1"
         or os.getenv("TORCH_CUDA_ARCH_LIST", "") != ""
     ):