pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 pythonOlder,
6 fsspec,
7 lightning-utilities,
8 numpy,
9 packaging,
10 pyyaml,
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.2.4";
25 format = "pyproject";
26
27 src = fetchFromGitHub {
28 owner = "Lightning-AI";
29 repo = "pytorch-lightning";
30 rev = "refs/tags/${version}";
31 hash = "sha256-IkoSID7nEPbKrhEMlo/UaMcF80HYldvndFA54DoHT+M=";
32 };
33
34 preConfigure = ''
35 export PACKAGE_NAME=pytorch
36 '';
37
38 propagatedBuildInputs = [
39 fsspec
40 numpy
41 packaging
42 pyyaml
43 tensorboardx
44 torch
45 lightning-utilities
46 torchmetrics
47 tqdm
48 traitlets
49 ] ++ fsspec.optional-dependencies.http;
50
51 nativeCheckInputs = [
52 psutil
53 pytestCheckHook
54 ];
55
56 # Some packages are not in NixPkgs; other tests try to build distributed
57 # models, which doesn't work in the sandbox.
58 doCheck = false;
59
60 pythonImportsCheck = [ "pytorch_lightning" ];
61
62 meta = with lib; {
63 description = "Lightweight PyTorch wrapper for machine learning researchers";
64 homepage = "https://pytorch-lightning.readthedocs.io";
65 license = licenses.asl20;
66 maintainers = with maintainers; [ tbenst ];
67 };
68}