8.6 requires CUDA 11.1 (#1289)

e1aa1fc1 · Masaki Kozuki · GitHub · a786ca0c · e1aa1fc1
Unverified Commit e1aa1fc1 authored Feb 10, 2022 by Masaki Kozuki Committed by GitHub Feb 10, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

setup.py setup.py +4 -2

No files found.
--- a/setup.py
+++ b/setup.py
@@ -72,9 +72,11 @@ if not torch.cuda.is_available():
        'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n',
    )
    if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
-        _, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
+        _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
        if int(bare_metal_major) == 11:
-            os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6"
+            os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0"
+            if int(bare_metal_minor) > 0:
+                os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6"
        else:
            os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"