diff options
| author | Adam J. Stewart <ajstewart426@gmail.com> | 2023-11-09 16:25:54 -0600 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-11-09 14:25:54 -0800 |
| commit | b8bb8a70cea7106a9eb71ad28815b00bdaf85edb (patch) | |
| tree | 71e48fc77631f0928a25482590138c9223a23a6d | |
| parent | dd2b436b5a6d37c13254f9d85cb25b4a67d7e21b (diff) | |
| download | spack-b8bb8a70cea7106a9eb71ad28815b00bdaf85edb.tar.gz spack-b8bb8a70cea7106a9eb71ad28815b00bdaf85edb.tar.bz2 spack-b8bb8a70cea7106a9eb71ad28815b00bdaf85edb.tar.xz spack-b8bb8a70cea7106a9eb71ad28815b00bdaf85edb.zip | |
PyTorch: specify CUDA root directory (#40855)
| -rw-r--r-- | var/spack/repos/builtin/packages/py-torch/package.py | 5 |
1 files changed, 2 insertions, 3 deletions
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py index d2edd94538..782b0741d2 100644 --- a/var/spack/repos/builtin/packages/py-torch/package.py +++ b/var/spack/repos/builtin/packages/py-torch/package.py @@ -491,9 +491,8 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage): 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 ) |