zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add additional cuda-toolkit location variable used by py-torch (#46245)

Browse Source 
superceds #46128

Co-authored-by: Robert Underwood <runderwood@anl.gov>
This commit is contained in:
Robert Underwood 2024-09-10 11:20:05 -04:00 committed by GitHub
parent 6f5ba44431
commit 79d778f8cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
var/spack/repos/builtin/packages/py-torch/package.py
View File

@ -526,6 +526,7 @@ def enable_or_disable(variant, keyword="USE", var=None):
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)