diff options
| -rw-r--r-- | var/spack/repos/builtin/packages/py-torch/package.py | 3 |
1 files changed, 1 insertions, 2 deletions
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py index b34d0b08f7..73f191d306 100644 --- a/var/spack/repos/builtin/packages/py-torch/package.py +++ b/var/spack/repos/builtin/packages/py-torch/package.py @@ -7,7 +7,7 @@ from spack import * # TODO: try switching to CMakePackage for more control over build -class PyTorch(PythonPackage): +class PyTorch(PythonPackage, CudaPackage): """Tensors and Dynamic neural networks in Python with strong GPU acceleration.""" @@ -61,7 +61,6 @@ class PyTorch(PythonPackage): version('0.4.0', tag='v0.4.0', submodules=True) version('0.3.1', tag='v0.3.1', submodules=True) - variant('cuda', default=True, description='Enables CUDA build') variant('cudnn', default=True, description='Enables the cuDNN build') variant('magma', default=False, description='Enables the MAGMA build') variant('fbgemm', default=False, description='Enables the FBGEMM build') |