tjh.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, ninja
5, which
6, libjpeg_turbo
7, libpng
8, numpy
9, scipy
10, pillow
11, pytorch
12, pytest
13}:
14
15buildPythonPackage rec {
16 pname = "torchvision";
17 version = "0.9.0";
18
19 src = fetchFromGitHub {
20 owner = "pytorch";
21 repo = "vision";
22 rev = "v${version}";
23 sha256 = "0rw186nki7q8igk481p2mvx4n763i3kpn663yp3dibkimddpnf9k";
24 };
25
26 nativeBuildInputs = [ libpng ninja which ];
27
28 TORCHVISION_INCLUDE = "${libjpeg_turbo.dev}/include/";
29 TORCHVISION_LIBRARY = "${libjpeg_turbo}/lib/";
30
31 buildInputs = [ libjpeg_turbo libpng ];
32
33 propagatedBuildInputs = [ numpy pillow pytorch scipy ];
34
35 # tries to download many datasets for tests
36 doCheck = false;
37
38 checkPhase = ''
39 HOME=$TMPDIR py.test test --ignore=test/test_datasets_download.py
40 '';
41
42 checkInputs = [ pytest ];
43
44 meta = with lib; {
45 description = "PyTorch vision library";
46 homepage = "https://pytorch.org/";
47 license = licenses.bsd3;
48 maintainers = with maintainers; [ ericsagnes ];
49 };
50}