Check for _is_cuda() in compute_num_jobs (#3481)

ba8ae1d8 · bnellnm · GitHub · 84eaa684 · ba8ae1d8
Unverified Commit ba8ae1d8 authored Mar 20, 2024 by bnellnm Committed by GitHub Mar 20, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 6 deletions

setup.py setup.py +6 -6

No files found.
--- a/setup.py
+++ b/setup.py
@@ -61,12 +61,12 @@ class cmake_build_ext(build_ext):
        except AttributeError:
            num_jobs = os.cpu_count()
-        nvcc_cuda_version = get_nvcc_cuda_version()
+        nvcc_threads = None
-        if nvcc_cuda_version >= Version("11.2"):
+        if _is_cuda():
-            nvcc_threads = int(os.getenv("NVCC_THREADS", 8))
+            nvcc_cuda_version = get_nvcc_cuda_version()
-            num_jobs = max(1, round(num_jobs / (nvcc_threads / 4)))
+            if nvcc_cuda_version >= Version("11.2"):
-        else:
+                nvcc_threads = int(os.getenv("NVCC_THREADS", 8))
-            nvcc_threads = None
+                num_jobs = max(1, round(num_jobs / (nvcc_threads / 4)))
        return num_jobs, nvcc_threads