about summary refs log tree commit diff
diff options
context:
space:
mode:
authorMatthew Bauer <mjbauer95@gmail.com>2020-09-10 14:23:15 -0500
committerMatthew Bauer <mjbauer95@gmail.com>2020-09-10 14:23:15 -0500
commitc85fc341b04548c3442ceb7027004f822f0a052a (patch)
treee0da47808e8913a2bbad854e3253f210e4e54341
parent34f475f5eae13d18b4e4b8a17aa7a772d8619b0b (diff)
pytorch: update to use passthru's for checks
This makes the checks a little easier to understand - we can just
check it in the passthru. Also increase major check for cuda 11.
-rw-r--r--pkgs/development/libraries/science/math/magma/default.nix2
-rw-r--r--pkgs/development/python-modules/pytorch/default.nix10
2 files changed, 5 insertions, 7 deletions
diff --git a/pkgs/development/libraries/science/math/magma/default.nix b/pkgs/development/libraries/science/math/magma/default.nix
index 38700c963bf81..2079ace021b68 100644
--- a/pkgs/development/libraries/science/math/magma/default.nix
+++ b/pkgs/development/libraries/science/math/magma/default.nix
@@ -50,4 +50,6 @@ in stdenv.mkDerivation {
     platforms = platforms.unix;
     maintainers = with maintainers; [ tbenst ];
   };
+
+  passthru.cudatoolkit = cudatoolkit;
 }
diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix
index 9bdead706d7b1..4635f813c3a6c 100644
--- a/pkgs/development/python-modules/pytorch/default.nix
+++ b/pkgs/development/python-modules/pytorch/default.nix
@@ -25,15 +25,11 @@ assert !openMPISupport || openmpi != null;
 assert !cudaSupport || cudatoolkit != null;
 assert cudnn == null || cudatoolkit != null;
 assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version;
-                        in majorIs == "9" || majorIs == "10");
+                        in majorIs == "9" || majorIs == "10" || majorIs == "11");
 
-let
-  hasDependency = dep: pkg: lib.lists.any (inp: inp == dep) pkg.buildInputs;
-  matchesCudatoolkit = hasDependency cudatoolkit;
-in
 # confirm that cudatoolkits are sync'd across dependencies
-assert !(openMPISupport && cudaSupport) || matchesCudatoolkit openmpi;
-assert !cudaSupport || matchesCudatoolkit magma;
+assert !(openMPISupport && cudaSupport) || openmpi.cudatoolkit == cudatoolkit;
+assert !cudaSupport || magma.cudatoolkit == cudatoolkit;
 
 let
   cudatoolkit_joined = symlinkJoin {