From db5aa927db12cb575e79c49b969fab92b85a2be3 Mon Sep 17 00:00:00 2001
From: Phil Tooley <32297355+ptooley@users.noreply.github.com>
Date: Tue, 10 Nov 2020 18:28:01 +0000
Subject: 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
---
 var/spack/repos/builtin/packages/py-torchvision/package.py | 5 +++++
 1 file changed, 5 insertions(+)

(limited to 'var')

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)
-- 
cgit v1.2.3-60-g2f50