diff options
Diffstat (limited to 'pkgs/development/python-modules/torchvision/default.nix')
-rw-r--r-- | pkgs/development/python-modules/torchvision/default.nix | 71 |
1 files changed, 40 insertions, 31 deletions
diff --git a/pkgs/development/python-modules/torchvision/default.nix b/pkgs/development/python-modules/torchvision/default.nix index 496a8c08710b0..91a269170e3bf 100644 --- a/pkgs/development/python-modules/torchvision/default.nix +++ b/pkgs/development/python-modules/torchvision/default.nix @@ -1,15 +1,16 @@ -{ buildPythonPackage -, fetchFromGitHub -, lib -, libjpeg_turbo -, libpng -, ninja -, numpy -, pillow -, pytest -, scipy -, torch -, which +{ + buildPythonPackage, + fetchFromGitHub, + lib, + libjpeg_turbo, + libpng, + ninja, + numpy, + pillow, + pytest, + scipy, + torch, + which, }: let @@ -17,7 +18,7 @@ let inherit (cudaPackages) backendStdenv; pname = "torchvision"; - version = "0.17.2"; + version = "0.18.0"; in buildPythonPackage { inherit pname version; @@ -26,33 +27,41 @@ buildPythonPackage { owner = "pytorch"; repo = "vision"; rev = "refs/tags/v${version}"; - hash = "sha256-Y7TkdbdrdiXQO2pBkUePCLw1EYQjsAqkV2OTMzxbU78="; + hash = "sha256-VWbalbLSV5a+t9eAO7QzQ/e11KkhGg6MHgd5vXcAUXc="; }; nativeBuildInputs = [ libpng ninja which - ] ++ lib.optionals cudaSupport [ - cudaPackages.cuda_nvcc - ]; + ] ++ lib.optionals cudaSupport [ cudaPackages.cuda_nvcc ]; - buildInputs = [ libjpeg_turbo libpng torch.cxxdev ]; + buildInputs = [ + libjpeg_turbo + libpng + torch.cxxdev + ]; - propagatedBuildInputs = [ numpy pillow torch scipy ]; + propagatedBuildInputs = [ + numpy + pillow + torch + scipy + ]; - preConfigure = '' - export TORCHVISION_INCLUDE="${libjpeg_turbo.dev}/include/" - export TORCHVISION_LIBRARY="${libjpeg_turbo}/lib/" - '' - # NOTE: We essentially override the compilers provided by stdenv because we don't have a hook - # for cudaPackages to swap in compilers supported by NVCC. - + lib.optionalString cudaSupport '' - export CC=${backendStdenv.cc}/bin/cc - export CXX=${backendStdenv.cc}/bin/c++ - export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}" - export FORCE_CUDA=1 - ''; + preConfigure = + '' + export TORCHVISION_INCLUDE="${libjpeg_turbo.dev}/include/" + export TORCHVISION_LIBRARY="${libjpeg_turbo}/lib/" + '' + # NOTE: We essentially override the compilers provided by stdenv because we don't have a hook + # for cudaPackages to swap in compilers supported by NVCC. + + lib.optionalString cudaSupport '' + export CC=${backendStdenv.cc}/bin/cc + export CXX=${backendStdenv.cc}/bin/c++ + export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}" + export FORCE_CUDA=1 + ''; # tries to download many datasets for tests doCheck = false; |