version guard (#1253)

e8473822 · Masaki Kozuki · GitHub · e0f5ea8c · e8473822
Unverified Commit e8473822 authored Dec 16, 2021 by Masaki Kozuki Committed by GitHub Dec 16, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

setup.py setup.py +4 -1

No files found.
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,10 @@ def raise_if_cuda_home_none(global_option: str) -> None:


 def append_nvcc_threads(nvcc_extra_args):
-    return nvcc_extra_args + ["--threads", "4"]
+    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
+    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
+        return nvcc_extra_args + ["--threads", "4"]
+    return nvcc_extra_args


 if not torch.cuda.is_available():