1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, isPy27
5, future
6, fsspec
7, packaging
8, pytestCheckHook
9, torch
10, pyyaml
11, tensorboard
12, torchmetrics
13, tqdm }:
14
15buildPythonPackage rec {
16 pname = "pytorch-lightning";
17 version = "1.6.5";
18
19 disabled = isPy27;
20
21 src = fetchFromGitHub {
22 owner = "PyTorchLightning";
23 repo = pname;
24 rev = "refs/tags/${version}";
25 hash = "sha256-CgD5g5nhz2DI4gOQyPl8/Cq6wWHzL0ALgOB5SgUOgaI=";
26 };
27
28 propagatedBuildInputs = [
29 packaging
30 future
31 fsspec
32 torch
33 pyyaml
34 tensorboard
35 torchmetrics
36 tqdm
37 ];
38
39 checkInputs = [ pytestCheckHook ];
40 # Some packages are not in NixPkgs; other tests try to build distributed
41 # models, which doesn't work in the sandbox.
42 doCheck = false;
43
44 pythonImportsCheck = [ "pytorch_lightning" ];
45
46 meta = with lib; {
47 description = "Lightweight PyTorch wrapper for machine learning researchers";
48 homepage = "https://pytorch-lightning.readthedocs.io";
49 license = licenses.asl20;
50 maintainers = with maintainers; [ tbenst ];
51 };
52}