diff options
| -rw-r--r-- | var/spack/repos/builtin/packages/py-tensorflow/package.py | 12 |
1 files changed, 8 insertions, 4 deletions
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/package.py b/var/spack/repos/builtin/packages/py-tensorflow/package.py index f139cf6cb3..69dabac970 100644 --- a/var/spack/repos/builtin/packages/py-tensorflow/package.py +++ b/var/spack/repos/builtin/packages/py-tensorflow/package.py @@ -205,6 +205,11 @@ class PyTensorflow(Package, CudaPackage): conflicts('+computecpp', when='~opencl') conflicts('+rocm', when='@:1.11') conflicts('+cuda', when='platform=darwin', msg='There is no GPU support for macOS') + conflicts('cuda_arch=none', when='+cuda', msg='Must specify CUDA compute capabilities of your GPU, see https://developer.nvidia.com/cuda-gpus') + conflicts('cuda_arch=20', when='@1.12.1,1.14:', msg='TensorFlow only supports compute capabilities >= 3.5') + conflicts('cuda_arch=30', when='@1.12.1,1.14:', msg='TensorFlow only supports compute capabilities >= 3.5') + conflicts('cuda_arch=32', when='@1.12.1,1.14:', msg='TensorFlow only supports compute capabilities >= 3.5') + conflicts('cuda_arch=20', when='@1.4:1.12.0,1.12.2:1.12.3', msg='Only compute capabilities 3.0 or higher are supported') conflicts('+tensorrt', when='@:1.5') conflicts('+tensorrt', when='~cuda') conflicts('+tensorrt', when='platform=darwin', msg='Currently TensorRT is only supported on Linux platform') @@ -423,10 +428,9 @@ class PyTensorflow(Package, CudaPackage): # Please note that each additional compute capability significantly # increases your build time and binary size, and that TensorFlow # only supports compute capabilities >= 3.5 - if spec.variants['cuda_arch'].value != 'none': - capabilities = ','.join('{0:.1f}'.format( - float(i) / 10.0) for i in spec.variants['cuda_arch'].value) - env.set('TF_CUDA_COMPUTE_CAPABILITIES', capabilities) + capabilities = ','.join('{0:.1f}'.format( + float(i) / 10.0) for i in spec.variants['cuda_arch'].value) + env.set('TF_CUDA_COMPUTE_CAPABILITIES', capabilities) else: env.set('TF_NEED_CUDA', '0') |