summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorRobert Underwood <robertu94@users.noreply.github.com>2024-09-10 11:20:05 -0400
committerGitHub <noreply@github.com>2024-09-10 17:20:05 +0200
commit79d778f8cdb5f4f5eaaeb9a5d1215080698041e3 (patch)
tree7afca6a6d3c1e749a921a40856d1c9d73fcee87a
parent6f5ba4443119cc9fb67e025adea32a92ba07d469 (diff)
downloadspack-79d778f8cdb5f4f5eaaeb9a5d1215080698041e3.tar.gz
spack-79d778f8cdb5f4f5eaaeb9a5d1215080698041e3.tar.bz2
spack-79d778f8cdb5f4f5eaaeb9a5d1215080698041e3.tar.xz
spack-79d778f8cdb5f4f5eaaeb9a5d1215080698041e3.zip
Add additional cuda-toolkit location variable used by py-torch (#46245)
superceds #46128 Co-authored-by: Robert Underwood <runderwood@anl.gov>
-rw-r--r--var/spack/repos/builtin/packages/py-torch/package.py1
1 files changed, 1 insertions, 0 deletions
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py
index 9af10cb65a..2e1631fa02 100644
--- a/var/spack/repos/builtin/packages/py-torch/package.py
+++ b/var/spack/repos/builtin/packages/py-torch/package.py
@@ -526,6 +526,7 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
enable_or_disable("cuda")
if "+cuda" in self.spec:
+ env.set("CUDA_TOOLKIT_ROOT_DIR", self.spec["cuda"].prefix) # Linux/macOS
env.set("CUDA_HOME", self.spec["cuda"].prefix) # Linux/macOS
env.set("CUDA_PATH", self.spec["cuda"].prefix) # Windows
self.torch_cuda_arch_list(env)