about summary refs log tree commit diff
path: root/nixpkgs/pkgs/development/python-modules/pytorch/default.nix
diff options
context:
space:
mode:
Diffstat (limited to 'nixpkgs/pkgs/development/python-modules/pytorch/default.nix')
-rw-r--r--nixpkgs/pkgs/development/python-modules/pytorch/default.nix10
1 files changed, 4 insertions, 6 deletions
diff --git a/nixpkgs/pkgs/development/python-modules/pytorch/default.nix b/nixpkgs/pkgs/development/python-modules/pytorch/default.nix
index f072972937a9..db1914f4ee7b 100644
--- a/nixpkgs/pkgs/development/python-modules/pytorch/default.nix
+++ b/nixpkgs/pkgs/development/python-modules/pytorch/default.nix
@@ -1,7 +1,7 @@
 { stdenv, lib, fetchFromGitHub, fetchpatch, buildPythonPackage, python,
   cudaSupport ? false, cudatoolkit ? null, cudnn ? null, nccl ? null, magma ? null,
   mklDnnSupport ? true, useSystemNccl ? true,
-  openMPISupport ? false, openmpi ? null,
+  MPISupport ? false, mpi,
   buildDocs ? false,
   cudaArchList ? null,
 
@@ -29,8 +29,6 @@
 
   isPy3k, pythonOlder }:
 
-assert !openMPISupport || openmpi != null;
-
 # assert that everything needed for cuda is present and that the correct cuda versions are used
 assert !cudaSupport || cudatoolkit != null;
 assert cudnn == null || cudatoolkit != null;
@@ -38,7 +36,7 @@ assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version;
                         in majorIs == "9" || majorIs == "10" || majorIs == "11");
 
 # confirm that cudatoolkits are sync'd across dependencies
-assert !(openMPISupport && cudaSupport) || openmpi.cudatoolkit == cudatoolkit;
+assert !(MPISupport && cudaSupport) || mpi.cudatoolkit == cudatoolkit;
 assert !cudaSupport || magma.cudatoolkit == cudatoolkit;
 
 let
@@ -224,7 +222,7 @@ in buildPythonPackage rec {
     typing-extensions
     # the following are required for tensorboard support
     pillow six future tensorflow-tensorboard protobuf
-  ] ++ lib.optionals openMPISupport [ openmpi ]
+  ] ++ lib.optionals MPISupport [ mpi ]
     ++ lib.optionals (pythonOlder "3.7") [ dataclasses ];
 
   checkInputs = [ hypothesis ninja psutil ];
@@ -267,7 +265,7 @@ in buildPythonPackage rec {
     cp -r $out/${python.sitePackages}/torch/lib     $lib/lib
   '';
 
-  postFixup = stdenv.lib.optionalString stdenv.isDarwin ''
+  postFixup = lib.optionalString stdenv.isDarwin ''
     for f in $(ls $lib/lib/*.dylib); do
         install_name_tool -id $lib/lib/$(basename $f) $f || true
     done