tjh.dev
1{ lib
2, stdenv
3, buildPythonPackage
4, fetchurl
5, isPy37
6, isPy38
7, isPy39
8, patchelf
9, pillow
10, python
11, pytorch-bin
12}:
13
14let
15 pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
16 srcs = import ./binary-hashes.nix version;
17 unsupported = throw "Unsupported system";
18 version = "0.9.1";
19in buildPythonPackage {
20 inherit version;
21
22 pname = "torchvision";
23
24 format = "wheel";
25
26 src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
27
28 disabled = !(isPy37 || isPy38 || isPy39);
29
30 nativeBuildInputs = [
31 patchelf
32 ];
33
34 propagatedBuildInputs = [
35 pillow
36 pytorch-bin
37 ];
38
39 pythonImportsCheck = [ "torchvision" ];
40
41 postFixup = let
42 rpath = lib.makeLibraryPath [ stdenv.cc.cc.lib ];
43 in ''
44 # Note: after patchelf'ing, libcudart can still not be found. However, this should
45 # not be an issue, because PyTorch is loaded before torchvision and brings
46 # in the necessary symbols.
47 patchelf --set-rpath "${rpath}:${pytorch-bin}/${python.sitePackages}/torch/lib:" \
48 "$out/${python.sitePackages}/torchvision/_C.so"
49 '';
50
51 meta = with lib; {
52 description = "PyTorch vision library";
53 homepage = "https://pytorch.org/";
54 changelog = "https://github.com/pytorch/vision/releases/tag/v${version}";
55 license = licenses.bsd3;
56 platforms = platforms.linux;
57 maintainers = with maintainers; [ danieldk ];
58 };
59}