summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--var/spack/repos/builtin/packages/py-torchvision/package.py1
1 files changed, 1 insertions, 0 deletions
diff --git a/var/spack/repos/builtin/packages/py-torchvision/package.py b/var/spack/repos/builtin/packages/py-torchvision/package.py
index ae6402d1c1..b70e2cd367 100644
--- a/var/spack/repos/builtin/packages/py-torchvision/package.py
+++ b/var/spack/repos/builtin/packages/py-torchvision/package.py
@@ -86,5 +86,6 @@ class PyTorchvision(PythonPackage):
if '+cuda' in self.spec['py-torch']:
env.set('FORCE_CUDA', 1)
+ env.set('CUDA_HOME', self.spec['cuda'].prefix)
else:
env.set('FORCE_CUDA', 0)