diff options
Diffstat (limited to 'pkgs/development/python-modules/torchinfo/default.nix')
| -rw-r--r-- | pkgs/development/python-modules/torchinfo/default.nix | 30 |
1 files changed, 25 insertions, 5 deletions
diff --git a/pkgs/development/python-modules/torchinfo/default.nix b/pkgs/development/python-modules/torchinfo/default.nix index 420a5fd8dfc52..59b7e99c9dfca 100644 --- a/pkgs/development/python-modules/torchinfo/default.nix +++ b/pkgs/development/python-modules/torchinfo/default.nix @@ -1,26 +1,38 @@ { lib , buildPythonPackage , fetchFromGitHub -, pytestCheckHook +, fetchpatch , pythonOlder , torch , torchvision +, pytestCheckHook +, transformers }: buildPythonPackage rec { pname = "torchinfo"; - version = "1.7.2"; + version = "1.8.0"; format = "setuptools"; disabled = pythonOlder "3.7"; src = fetchFromGitHub { owner = "TylerYep"; - repo = pname; + repo = "torchinfo"; rev = "refs/tags/v${version}"; - hash = "sha256-O+I7BNQ5moV/ZcbbuP/IFoi0LO0WsGHBbSfgPmFu1Ec="; + hash = "sha256-pPjg498aT8y4b4tqIzNxxKyobZX01u+66ScS/mee51Q="; }; + patches = [ + (fetchpatch { # Add support for Python 3.11 and pytorch 2.1 + url = "https://github.com/TylerYep/torchinfo/commit/c74784c71c84e62bcf56664653b7f28d72a2ee0d.patch"; + hash = "sha256-xSSqs0tuFpdMXUsoVv4sZLCeVnkK6pDDhX/Eobvn5mw="; + includes = [ + "torchinfo/model_statistics.py" + ]; + }) + ]; + propagatedBuildInputs = [ torch torchvision @@ -28,17 +40,25 @@ buildPythonPackage rec { nativeCheckInputs = [ pytestCheckHook + transformers ]; + preCheck = '' + export HOME=$(mktemp -d) + ''; + disabledTests = [ # Skip as it downloads pretrained weights (require network access) "test_eval_order_doesnt_matter" + "test_flan_t5_small" # AssertionError in output "test_google" + # "addmm_impl_cpu_" not implemented for 'Half' + "test_input_size_half_precision" ]; disabledTestPaths = [ - # Wants "compressai", which we don't package (2023-03-23) + # Test requires network access "tests/torchinfo_xl_test.py" ]; |