diff options
Diffstat (limited to 'pkgs/development/python-modules/pytorch/default.nix')
| -rw-r--r-- | pkgs/development/python-modules/pytorch/default.nix | 14 |
1 files changed, 8 insertions, 6 deletions
diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix index 145c9a240750..1436153e1db5 100644 --- a/pkgs/development/python-modules/pytorch/default.nix +++ b/pkgs/development/python-modules/pytorch/default.nix @@ -38,6 +38,7 @@ assert !(MPISupport && cudaSupport) || mpi.cudatoolkit == cudatoolkit; assert !cudaSupport || magma.cudatoolkit == cudatoolkit; let + setBool = v: if v then "1" else "0"; cudatoolkit_joined = symlinkJoin { name = "${cudatoolkit.name}-unsplit"; # nccl is here purely for semantic grouping it could be moved to nativeBuildInputs @@ -160,16 +161,17 @@ in buildPythonPackage rec { # Use pytorch's custom configurations dontUseCmakeConfigure = true; - BUILD_NAMEDTENSOR = true; - BUILD_DOCS = buildDocs; + BUILD_NAMEDTENSOR = setBool true; + BUILD_DOCS = setBool buildDocs; - USE_MKL = blas.implementation == "mkl"; + # We only do an imports check, so do not build tests either. + BUILD_TEST = setBool false; # Unlike MKL, oneDNN (née MKLDNN) is FOSS, so we enable support for # it by default. PyTorch currently uses its own vendored version # of oneDNN through Intel iDeep. - USE_MKLDNN = mklDnnSupport; - USE_MKLDNN_CBLAS = mklDnnSupport; + USE_MKLDNN = setBool mklDnnSupport; + USE_MKLDNN_CBLAS = setBool mklDnnSupport; preBuild = '' export MAX_JOBS=$NIX_BUILD_CORES @@ -198,7 +200,7 @@ in buildPythonPackage rec { PYTORCH_BUILD_VERSION = version; PYTORCH_BUILD_NUMBER = 0; - USE_SYSTEM_NCCL=useSystemNccl; # don't build pytorch's third_party NCCL + USE_SYSTEM_NCCL=setBool useSystemNccl; # don't build pytorch's third_party NCCL # Suppress a weird warning in mkl-dnn, part of ideep in pytorch # (upstream seems to have fixed this in the wrong place?) |