diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py
index dcec15dc68..2a20235bef 100644
--- a/var/spack/repos/builtin/packages/py-torch/package.py
+++ b/var/spack/repos/builtin/packages/py-torch/package.py
@@ -62,6 +62,7 @@ class PyTorch(PythonPackage, CudaPackage):
     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='Build with CUDA')
     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')
@@ -100,6 +101,27 @@ class PyTorch(PythonPackage, CudaPackage):
     conflicts('+zstd', when='@:1.0')
     conflicts('+tbb', when='@:1.1')
 
+    cuda_arch_conflict = ('This version of Torch/Caffe2 only supports compute '
+                          'capabilities ')
+
+    conflicts('cuda_arch=none', when='+cuda+caffe2',
+              msg='Must specify CUDA compute capabilities of your GPU, see '
+              'https://developer.nvidia.com/cuda-gpus')
+    conflicts('cuda_arch=52', when='@1.3.0:+cuda+caffe2',
+              msg=cuda_arch_conflict + '>=5.3')
+    conflicts('cuda_arch=50', when='@1.3.0:+cuda+caffe2',
+              msg=cuda_arch_conflict + '>=5.3')
+    conflicts('cuda_arch=35', when='@1.3.0:+cuda+caffe2',
+              msg=cuda_arch_conflict + '>=5.3')
+    conflicts('cuda_arch=32', when='@1.3.0:+cuda+caffe2',
+              msg=cuda_arch_conflict + '>=5.3')
+    conflicts('cuda_arch=30', when='@1.3.0:+cuda+caffe2',
+              msg=cuda_arch_conflict + '>=5.3')
+    conflicts('cuda_arch=30', when='@1.2.0:+cuda+caffe2',
+              msg=cuda_arch_conflict + '>=3.2')
+    conflicts('cuda_arch=20', when='@1.0.0:+cuda+caffe2',
+              msg=cuda_arch_conflict + '>=3.0')
+
     # Required dependencies
     depends_on('cmake@3.5:', type='build')
     # Use Ninja generator to speed up build times
@@ -128,7 +150,10 @@ class PyTorch(PythonPackage, CudaPackage):
     # depends_on('fbgemm', when='+fbgemm')
     # TODO: add dependency: https://github.com/ROCmSoftwarePlatform/MIOpen
     # depends_on('miopen', when='+miopen')
-    depends_on('intel-mkl-dnn', when='+mkldnn')
+    # TODO: See if there is a way to use an external mkldnn installation.
+    # Currently, only older versions of py-torch use an external mkldnn
+    # library.
+    depends_on('intel-mkl-dnn', when='@0.4:0.4.1+mkldnn')
     # TODO: add dependency: https://github.com/Maratyszcza/NNPACK
     # depends_on('nnpack', when='+nnpack')
     depends_on('qnnpack', when='+qnnpack')
@@ -197,6 +222,10 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
         enable_or_disable('cuda')
         if '+cuda' in self.spec:
             env.set('CUDA_HOME', self.spec['cuda'].prefix)
+            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)
 
         enable_or_disable('cudnn')
         if '+cudnn' in self.spec:
@@ -213,7 +242,7 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
             env.set('MIOPEN_LIBRARY', self.spec['miopen'].libs[0])
 
         enable_or_disable('mkldnn')
-        if '+mkldnn' in self.spec:
+        if '@0.4:0.4.1+mkldnn' in self.spec:
             env.set('MKLDNN_HOME', self.spec['intel-mkl-dnn'].prefix)
 
         enable_or_disable('nnpack')