zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 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
1 changed files with 2 additions and 3 deletions
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") enable_or_disable("cuda")
if "+cuda" in self.spec: if "+cuda" in self.spec:
# cmake/public/cuda.cmake env.set("CUDA_HOME", self.spec["cuda"].prefix) # Linux/macOS
# cmake/Modules_CUDA_fix/upstream/FindCUDA.cmake env.set("CUDA_PATH", self.spec["cuda"].prefix) # Windows
env.unset("CUDA_ROOT")
torch_cuda_arch = ";".join( torch_cuda_arch = ";".join(
"{0:.1f}".format(float(i) / 10.0) for i in self.spec.variants["cuda_arch"].value "{0:.1f}".format(float(i) / 10.0) for i in self.spec.variants["cuda_arch"].value
) )