diff options
| -rw-r--r-- | var/spack/repos/builtin/packages/caffe/package.py | 10 |
1 files changed, 8 insertions, 2 deletions
diff --git a/var/spack/repos/builtin/packages/caffe/package.py b/var/spack/repos/builtin/packages/caffe/package.py index 1dab412e95..b2f54a3d2d 100644 --- a/var/spack/repos/builtin/packages/caffe/package.py +++ b/var/spack/repos/builtin/packages/caffe/package.py @@ -6,7 +6,7 @@ from spack import * -class Caffe(CMakePackage): +class Caffe(CMakePackage, CudaPackage): """Caffe is a deep learning framework made with expression, speed, and modularity in mind. It is developed by the Berkeley Vision and Learning Center (BVLC) and by community contributors.""" @@ -43,7 +43,7 @@ class Caffe(CMakePackage): depends_on('hdf5 +hl +cxx') # Optional dependencies - depends_on('opencv@3.2.0+core+highgui+imgproc', when='+opencv') + depends_on('opencv@3.2.0:3.4.12+core+highgui+imgproc+imgcodecs', when='+opencv') depends_on('leveldb', when='+leveldb') depends_on('lmdb', when='+lmdb') depends_on('python@2.7:', when='+python') @@ -87,4 +87,10 @@ class Caffe(CMakePackage): '-DCMAKE_CXX_COMPILER={0}'.format(self.spec['mpi'].mpicxx) ]) + if '+cuda' in spec: + if spec.variants['cuda_arch'].value[0] != 'none': + cuda_arch = spec.variants['cuda_arch'].value + args.append(self.define('CUDA_ARCH_NAME', 'Manual')) + args.append(self.define('CUDA_ARCH_BIN', ' '.join(cuda_arch))) + return args |