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.
This commit is contained in:
Matthew Bauer 2020-09-10 14:23:15 -05:00
parent 34f475f5ea
commit c85fc341b0
2 changed files with 5 additions and 7 deletions

View file

@ -50,4 +50,6 @@ in stdenv.mkDerivation {
platforms = platforms.unix;
maintainers = with maintainers; [ tbenst ];
};
passthru.cudatoolkit = cudatoolkit;
}

View file

@ -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 {