tjh.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pytest
5, matplotlib
6, mock
7, pytorch
8, pynvml
9, scikitlearn
10, tqdm
11}:
12
13buildPythonPackage rec {
14 pname = "ignite";
15 version = "0.3.0";
16
17 src = fetchFromGitHub {
18 owner = "pytorch";
19 repo = pname;
20 rev = "v${version}";
21 sha256 = "0i863kxi1r1hspj19lhn6r8256vdazjcyvis0s33fgzrf7kxi08x";
22 };
23
24 checkInputs = [ pytest matplotlib mock ];
25 propagatedBuildInputs = [ pytorch scikitlearn tqdm pynvml ];
26
27 # Some packages are not in NixPkgs; other tests try to build distributed
28 # models, which doesn't work in the sandbox.
29 checkPhase = ''
30 pytest -k 'not visdom and not tensorboard and not mlflow and not polyaxon and not conftest and not engines and not distrib_' tests/
31 '';
32
33 meta = with lib; {
34 description = "High-level training library for PyTorch";
35 homepage = "https://pytorch.org/ignite";
36 license = licenses.bsd3;
37 maintainers = [ maintainers.bcdarwin ];
38 };
39}