tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{ lib
2, fetchPypi
3, python
4, buildPythonPackage
5, pythonOlder
6, pytorch
7, pytestCheckHook
8, torchvision
9}:
10
11buildPythonPackage rec {
12 pname = "torchinfo";
13 version = "1.6.5";
14 disabled = pythonOlder "3.7";
15
16 src = fetchPypi {
17 inherit pname version;
18 sha256 = "sha256-Vg/TXD+/VMIv1wHywaOuEj4MDTq90lUo99n+Nppu0uI=";
19 };
20
21 propagatedBuildInputs = [
22 pytorch
23 torchvision
24 ];
25
26 checkInputs = [
27 pytestCheckHook
28 ];
29
30 disabledTests = [
31 # Skip as it downloads pretrained weights (require network access)
32 "test_eval_order_doesnt_matter"
33 ];
34
35 pythonImportsCheck = [ "torchvision" ];
36
37 meta = {
38 description = "API to visualize pytorch models";
39 homepage = "https://github.com/TylerYep/torchinfo";
40 license = lib.licenses.mit;
41 maintainers = with lib.maintainers; [ petterstorvik ];
42 };
43}