diff options
Diffstat (limited to 'pkgs/development/python-modules/torchinfo/default.nix')
-rw-r--r-- | pkgs/development/python-modules/torchinfo/default.nix | 43 |
1 files changed, 43 insertions, 0 deletions
diff --git a/pkgs/development/python-modules/torchinfo/default.nix b/pkgs/development/python-modules/torchinfo/default.nix new file mode 100644 index 0000000000000..331729e25c581 --- /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.3"; + disabled = pythonOlder "3.7"; + + src = fetchPypi { + inherit pname version; + sha256 = "sha256-g1xhtdwygzPTTswP8iZ364ynBQE7D+aAsZ3d9EpyvIA="; + }; + + 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 ]; + }; +} |