diff --git a/pkgs/development/python-modules/torchinfo/default.nix b/pkgs/development/python-modules/torchinfo/default.nix new file mode 100644 index 000000000000..b9071782b590 --- /dev/null +++ b/pkgs/development/python-modules/torchinfo/default.nix @@ -0,0 +1,43 @@ +{ lib +, fetchPypi +, python +, buildPythonPackage +, pythonOlder +, pytorch +, pytestCheckHook +, torchvision +}: + +buildPythonPackage rec { + pname = "torchinfo"; + version = "1.6.2"; + disabled = pythonOlder "3.7"; + + src = fetchPypi { + inherit pname version; + sha256 = "sha256-vf3TADqIX4RtCU3dN/lBK4aRg3wud/KkK9u5XGnBbO4="; + }; + + propagatedBuildInputs = [ + pytorch + torchvision + ]; + + checkInputs = [ + pytestCheckHook + ]; + + disabledTests = [ + # Skip as it downloads pretrained weights (require network access) + "test_eval_order_doesnt_matter" + ]; + + pythonImportsCheck = [ "torchvision" ]; + + meta = { + description = "API to visualize pytorch models"; + homepage = "https://github.com/TylerYep/torchinfo"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ petterstorvik ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index f82a0e0851c9..715d828cc881 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -9728,6 +9728,8 @@ in { torchgpipe = callPackage ../development/python-modules/torchgpipe { }; + torchinfo = callPackage ../development/python-modules/torchinfo { }; + torchvision = callPackage ../development/python-modules/torchvision { }; torchvision-bin = callPackage ../development/python-modules/torchvision/bin.nix { };