Look for nvcc in CUDA_HOME

Author: amiller27

bindings/torch/setup.py

@@ -1,5 +1,6 @@
 import os
 
+from pathlib import Path
 import re
 from setuptools import setup
 from pkg_resources import parse_version
@@ -8,7 +9,7 @@
 import sys
 import torch
 from glob import glob
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
 
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
 ROOT_DIR = os.path.dirname(os.path.dirname(SCRIPT_DIR))
@@ -80,9 +81,14 @@ def find_cl_path():
 cpp_standard = 14
 
 # Get CUDA version and make sure the targeted compute capability is compatible
-if os.system("nvcc --version") == 0:
-	nvcc_out = subprocess.check_output(["nvcc", "--version"]).decode()
-	cuda_version = re.search(r"release (\S+),", nvcc_out)
+nvcc_version_result = subprocess.run(
+	[str(Path(CUDA_HOME) / "bin" / "nvcc"), "--version"],
+	text=True,
+	check=False,
+	stdout=subprocess.PIPE,
+)
+if nvcc_version_result.returncode == 0:
+	cuda_version = re.search(r"release (\S+),", nvcc_version_result.stdout)
 
 	if cuda_version:
 		cuda_version = parse_version(cuda_version.group(1))