tjh.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 fsspec,
6 lightning-utilities,
7 numpy,
8 packaging,
9 pyyaml,
10 setuptools,
11 tensorboardx,
12 torch,
13 torchmetrics,
14 tqdm,
15 traitlets,
16
17 # tests
18 psutil,
19 pytestCheckHook,
20}:
21
22buildPythonPackage rec {
23 pname = "pytorch-lightning";
24 version = "2.3.3";
25 pyproject = true;
26
27 src = fetchFromGitHub {
28 owner = "Lightning-AI";
29 repo = "pytorch-lightning";
30 rev = "refs/tags/${version}";
31 hash = "sha256-JgT+OiBToBBa3h556gF+kmcqLKLaK17tlTPokvEDyUE=";
32 };
33
34 preConfigure = ''
35 export PACKAGE_NAME=pytorch
36 '';
37
38 build-system = [ setuptools ];
39
40 dependencies = [
41 fsspec
42 numpy
43 packaging
44 pyyaml
45 tensorboardx
46 torch
47 lightning-utilities
48 torchmetrics
49 tqdm
50 traitlets
51 ] ++ fsspec.optional-dependencies.http;
52
53 nativeCheckInputs = [
54 psutil
55 pytestCheckHook
56 ];
57
58 # Some packages are not in NixPkgs; other tests try to build distributed
59 # models, which doesn't work in the sandbox.
60 doCheck = false;
61
62 pythonImportsCheck = [ "pytorch_lightning" ];
63
64 meta = with lib; {
65 description = "Lightweight PyTorch wrapper for machine learning researchers";
66 homepage = "https://github.com/Lightning-AI/pytorch-lightning";
67 changelog = "https://github.com/Lightning-AI/pytorch-lightning/releases/tag/${src.rev}";
68 license = licenses.asl20;
69 maintainers = with maintainers; [ tbenst ];
70 };
71}