pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pythonOlder
5, fsspec
6, lightning-utilities
7, numpy
8, packaging
9, pyyaml
10, tensorboardx
11, torch
12, torchmetrics
13, tqdm
14, traitlets
15
16# tests
17, psutil
18, pytestCheckHook
19}:
20
21buildPythonPackage rec {
22 pname = "pytorch-lightning";
23 version = "2.1.1";
24 format = "pyproject";
25
26 src = fetchFromGitHub {
27 owner = "Lightning-AI";
28 repo = "pytorch-lightning";
29 rev = "refs/tags/${version}";
30 hash = "sha256-1psTa++qF5WPDVXeDGWfcQ4hGz98uW297QDUKrQyoRE=";
31 };
32
33 preConfigure = ''
34 export PACKAGE_NAME=pytorch
35 '';
36
37 propagatedBuildInputs = [
38 fsspec
39 numpy
40 packaging
41 pyyaml
42 tensorboardx
43 torch
44 lightning-utilities
45 torchmetrics
46 tqdm
47 traitlets
48 ]
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 = [
61 "pytorch_lightning"
62 ];
63
64 meta = with lib; {
65 description = "Lightweight PyTorch wrapper for machine learning researchers";
66 homepage = "https://pytorch-lightning.readthedocs.io";
67 license = licenses.asl20;
68 maintainers = with maintainers; [ tbenst ];
69 };
70}