hlrs-software-stack/spack
7
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

PyTorch: specify CUDA root directory (#40855)

Browse source
This commit is contained in:
Adam J. Stewart 2023-11-09 16:25:54 -06:00 committed by GitHub
parent dd2b436b5a
commit b8bb8a70ce
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
var/spack/repos/builtin/packages/py-torch/package.py
View file

@ -491,9 +491,8 @@ def enable_or_disable(variant, keyword="USE", var=None, newer=False):
enable_or_disable("cuda")
if "+cuda" in self.spec:
# cmake/public/cuda.cmake
# cmake/Modules_CUDA_fix/upstream/FindCUDA.cmake
env.unset("CUDA_ROOT")
env.set("CUDA_HOME", self.spec["cuda"].prefix) # Linux/macOS
env.set("CUDA_PATH", self.spec["cuda"].prefix) # Windows
torch_cuda_arch = ";".join(
"{0:.1f}".format(float(i) / 10.0) for i in self.spec.variants["cuda_arch"].value
)