zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

py-torchvision: gather correct cuda_arch setting from py-torch (#19845)

Browse Source 
* py-torchvision: gather correct cuda_arch setting from py-torch dependency

* py-torchvision: remove unneeded variant
This commit is contained in:
Phil Tooley 2020-11-10 18:28:01 +00:00 committed by GitHub
parent f7f4171d4e
commit db5aa927db
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
var/spack/repos/builtin/packages/py-torchvision/package.py
View File

@ -98,5 +98,10 @@ def setup_build_environment(self, env):
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)