diff options
| author | Phil Tooley <32297355+ptooley@users.noreply.github.com> | 2020-11-10 18:28:01 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2020-11-10 12:28:01 -0600 |
| commit | db5aa927db12cb575e79c49b969fab92b85a2be3 (patch) | |
| tree | a329551ae8927aae692753b25736564c8bf36a65 /var | |
| parent | f7f4171d4ef8680fc558a4ef8e8c6e9af73d3ba9 (diff) | |
| download | spack-db5aa927db12cb575e79c49b969fab92b85a2be3.tar.gz spack-db5aa927db12cb575e79c49b969fab92b85a2be3.tar.bz2 spack-db5aa927db12cb575e79c49b969fab92b85a2be3.tar.xz spack-db5aa927db12cb575e79c49b969fab92b85a2be3.zip | |
py-torchvision: gather correct cuda_arch setting from py-torch (#19845)
* py-torchvision: gather correct cuda_arch setting from py-torch dependency
* py-torchvision: remove unneeded variant
Diffstat (limited to 'var')
| -rw-r--r-- | var/spack/repos/builtin/packages/py-torchvision/package.py | 5 |
1 files changed, 5 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 c358931453..aaca507ba9 100644 --- a/var/spack/repos/builtin/packages/py-torchvision/package.py +++ b/var/spack/repos/builtin/packages/py-torchvision/package.py @@ -98,5 +98,10 @@ class PyTorchvision(PythonPackage): if '+cuda' in self.spec['py-torch']: env.set('FORCE_CUDA', 1) env.set('CUDA_HOME', self.spec['cuda'].prefix) + pytorch_cuda_arch = ';'.join( + '{0:.1f}'.format(float(i) / 10.0) for i in + self.spec['py-torch'].variants['cuda_arch'].value + ) + env.set('TORCH_CUDA_ARCH_LIST', pytorch_cuda_arch) else: env.set('FORCE_CUDA', 0) |