diff --git a/pkgs/development/python-modules/torchtnt/default.nix b/pkgs/development/python-modules/torchtnt/default.nix new file mode 100644 index 000000000000..78c948823c77 --- /dev/null +++ b/pkgs/development/python-modules/torchtnt/default.nix @@ -0,0 +1,71 @@ +{ + lib, + buildPythonPackage, + fetchPypi, + + # build-system + setuptools, + + # dependencies + fsspec, + numpy, + packaging, + psutil, + pyre-extensions, + tabulate, + tensorboard, + torch, + tqdm, + typing-extensions, +}: + +buildPythonPackage rec { + pname = "torchtnt"; + version = "0.2.4"; + pyproject = true; + + # no tag / releases on github + src = fetchPypi { + inherit pname version; + hash = "sha256-Js9OcYllr8KT52FYtHKDciBVvPeelNDmfnC12/YcDJs="; + }; + + # requirements.txt is not included in Pypi archive + postPatch = '' + substituteInPlace setup.py \ + --replace-fail 'read_requirements("requirements.txt")' "[]" \ + --replace-fail 'read_requirements("dev-requirements.txt")' "[]" + ''; + + build-system = [ + setuptools + ]; + + dependencies = [ + fsspec + numpy + packaging + psutil + pyre-extensions + setuptools + tabulate + tensorboard + torch + tqdm + typing-extensions + ]; + + pythonImportsCheck = [ + "torchtnt" + ]; + + # Tests are not included in Pypi archive + doCheck = false; + + meta = { + description = "lightweight library for PyTorch training tools and utilities"; + homepage = "https://github.com/pytorch/tnt"; + license = lib.licenses.bsd3; + maintainers = with lib.maintainers; [ nim65s ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index ba43430e48c1..fd8b6c4a5ba6 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -16152,6 +16152,8 @@ self: super: with self; { torchsummary = callPackage ../development/python-modules/torchsummary { }; + torchtnt = callPackage ../development/python-modules/torchtnt { }; + torchvision = callPackage ../development/python-modules/torchvision { }; torchvision-bin = callPackage ../development/python-modules/torchvision/bin.nix {