pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pytestCheckHook
5, pythonOlder
6, torch
7, torchvision
8}:
9
10buildPythonPackage rec {
11 pname = "torchinfo";
12 version = "1.64";
13 format = "setuptools";
14
15 disabled = pythonOlder "3.7";
16
17 src = fetchFromGitHub {
18 owner = "TylerYep";
19 repo = pname;
20 rev = "refs/tags/v${version}";
21 hash = "sha256-gcl8RxCD017FP4LtB60WVtOh7jg2Otv/vNd9hKneEAU=";
22 };
23
24 propagatedBuildInputs = [
25 torch
26 torchvision
27 ];
28
29 checkInputs = [
30 pytestCheckHook
31 ];
32
33 disabledTests = [
34 # Skip as it downloads pretrained weights (require network access)
35 "test_eval_order_doesnt_matter"
36 # AssertionError in output
37 "test_google"
38 ];
39
40 pythonImportsCheck = [
41 "torchvision"
42 ];
43
44 meta = with lib; {
45 description = "API to visualize pytorch models";
46 homepage = "https://github.com/TylerYep/torchinfo";
47 license = licenses.mit;
48 maintainers = with maintainers; [ petterstorvik ];
49 };
50}