diff --git a/pkgs/development/python-modules/torchprofile/default.nix b/pkgs/development/python-modules/torchprofile/default.nix new file mode 100644 index 0000000000000..6a9cde7d6c34b --- /dev/null +++ b/pkgs/development/python-modules/torchprofile/default.nix @@ -0,0 +1,47 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + setuptools, + numpy, + torch, + torchvision, +}: + +buildPythonPackage rec { + pname = "torchprofile"; + version = "0.0.3"; + pyproject = true; + + src = fetchFromGitHub { + owner = "zhijian-liu"; + repo = "torchprofile"; + rev = "v${version}"; + hash = "sha256-6vxZHQwBjKpy288wcANdJ9gmvIOZloLv+iN76TtqYAI="; + }; + + build-system = [ + setuptools + ]; + + pythonRelaxDeps = [ + "torchvision" + ]; + + dependencies = [ + numpy + torch + torchvision + ]; + + pythonImportsCheck = [ + "torchprofile" + ]; + + meta = { + description = "A general and accurate MACs / FLOPs profiler for PyTorch models"; + homepage = "https://github.com/zhijian-liu/torchprofile"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ drupol ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 48bdda41d3ca6..dec0803ece86a 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -15921,6 +15921,8 @@ self: super: with self; { # Used by streamlit, 2021-01-29 tornado_5 = callPackage ../development/python-modules/tornado/5.nix { }; + torchprofile = callPackage ../development/python-modules/torchprofile { }; + torpy = callPackage ../development/python-modules/torpy { }; torrent-parser = callPackage ../development/python-modules/torrent-parser { };