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; platforms = platforms.unix;
maintainers = with maintainers; [ tbenst ]; maintainers = with maintainers; [ tbenst ];
}; };
passthru.cudatoolkit = cudatoolkit;
} }

View file

@ -25,15 +25,11 @@ assert !openMPISupport || openmpi != null;
assert !cudaSupport || cudatoolkit != null; assert !cudaSupport || cudatoolkit != null;
assert cudnn == null || cudatoolkit != null; assert cudnn == null || cudatoolkit != null;
assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version; 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 # confirm that cudatoolkits are sync'd across dependencies
assert !(openMPISupport && cudaSupport) || matchesCudatoolkit openmpi; assert !(openMPISupport && cudaSupport) || openmpi.cudatoolkit == cudatoolkit;
assert !cudaSupport || matchesCudatoolkit magma; assert !cudaSupport || magma.cudatoolkit == cudatoolkit;
let let
cudatoolkit_joined = symlinkJoin { cudatoolkit_joined = symlinkJoin {