summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--var/spack/repos/builtin/packages/py-torch/package.py7
1 files changed, 4 insertions, 3 deletions
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py
index 024e0d95b2..ce69a5581c 100644
--- a/var/spack/repos/builtin/packages/py-torch/package.py
+++ b/var/spack/repos/builtin/packages/py-torch/package.py
@@ -247,7 +247,7 @@ class PyTorch(PythonPackage, CudaPackage):
enable_or_disable('mkldnn')
if '@0.4:0.4.1+mkldnn' in self.spec:
- env.set('MKLDNN_HOME', self.spec['intel-mkl-dnn'].prefix)
+ env.set('MKLDNN_HOME', self.spec['onednn'].prefix)
enable_or_disable('nnpack')
enable_or_disable('qnnpack')
@@ -270,8 +270,9 @@ class PyTorch(PythonPackage, CudaPackage):
enable_or_disable('lmdb', newer=True)
enable_or_disable('binary', keyword='BUILD', newer=True)
- env.set('PYTORCH_BUILD_VERSION', self.version)
- env.set('PYTORCH_BUILD_NUMBER', 0)
+ if not self.spec.satisfies('@master'):
+ env.set('PYTORCH_BUILD_VERSION', self.version)
+ env.set('PYTORCH_BUILD_NUMBER', 0)
# BLAS to be used by Caffe2
if '^mkl' in self.spec: