diff options
Diffstat (limited to 'var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py')
| -rw-r--r-- | var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py | 44 |
1 files changed, 22 insertions, 22 deletions
diff --git a/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py b/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py index 81135e265e..ca91302ed7 100644 --- a/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py +++ b/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py @@ -9,36 +9,36 @@ from spack.package import * class PyTorchNvidiaApex(PythonPackage, CudaPackage): """A PyTorch Extension: Tools for easy mixed precision and - distributed training in Pytorch """ + distributed training in Pytorch""" homepage = "https://github.com/nvidia/apex/" - git = "https://github.com/nvidia/apex/" + git = "https://github.com/nvidia/apex/" - version('master', branch='master') - version('2020-10-19', commit='8a1ed9e8d35dfad26fb973996319965e4224dcdd') + version("master", branch="master") + version("2020-10-19", commit="8a1ed9e8d35dfad26fb973996319965e4224dcdd") - depends_on('python@3:', type=('build', 'run')) - depends_on('py-setuptools', type='build') - depends_on('py-torch@0.4:', type=('build', 'run')) - depends_on('cuda@9:', when='+cuda') - depends_on('py-pybind11', type=('build', 'link', 'run')) + depends_on("python@3:", type=("build", "run")) + depends_on("py-setuptools", type="build") + depends_on("py-torch@0.4:", type=("build", "run")) + depends_on("cuda@9:", when="+cuda") + depends_on("py-pybind11", type=("build", "link", "run")) - variant('cuda', default=True, description='Build with CUDA') + variant("cuda", default=True, description="Build with CUDA") def setup_build_environment(self, env): - if '+cuda' in self.spec: - env.set('CUDA_HOME', self.spec['cuda'].prefix) - if (self.spec.variants['cuda_arch'].value[0] != 'none'): - torch_cuda_arch = ';'.join( - '{0:.1f}'.format(float(i) / 10.0) for i - in - self.spec.variants['cuda_arch'].value) - env.set('TORCH_CUDA_ARCH_LIST', torch_cuda_arch) + if "+cuda" in self.spec: + env.set("CUDA_HOME", self.spec["cuda"].prefix) + if self.spec.variants["cuda_arch"].value[0] != "none": + torch_cuda_arch = ";".join( + "{0:.1f}".format(float(i) / 10.0) + for i in self.spec.variants["cuda_arch"].value + ) + env.set("TORCH_CUDA_ARCH_LIST", torch_cuda_arch) def install_options(self, spec, prefix): args = [] - if spec.satisfies('^py-torch@1.0:'): - args.append('--cpp_ext') - if '+cuda' in spec: - args.append('--cuda_ext') + if spec.satisfies("^py-torch@1.0:"): + args.append("--cpp_ext") + if "+cuda" in spec: + args.append("--cuda_ext") return args |