diff --git a/pkgs/development/python-modules/torch-ema/default.nix b/pkgs/development/python-modules/torch-ema/default.nix new file mode 100644 index 000000000000..a81ee06f9a19 --- /dev/null +++ b/pkgs/development/python-modules/torch-ema/default.nix @@ -0,0 +1,58 @@ +{ + lib, + stdenv, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + + # dependencies + torch, + + # tests + pytestCheckHook, +}: + +buildPythonPackage (finalAttrs: { + pname = "torch-ema"; + version = "0.3.0"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "fadel"; + repo = "pytorch_ema"; + tag = "v${finalAttrs.version}"; + hash = "sha256-OOF5Lb3HEIBXc1WXoUp7y44pheDc5oX/7L1vTrwNS2o="; + }; + + build-system = [ + setuptools + ]; + + dependencies = [ + torch + ]; + + pythonImportsCheck = [ "torch_ema" ]; + + nativeCheckInputs = [ + pytestCheckHook + ]; + + disabledTests = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [ + # aarch64-linux fails cpuinfo test, because /sys/devices/system/cpu/ does not exist in the sandbox: + # RuntimeError: Failed to initialize cpuinfo! + "test_state_dict_types" + "test_to" + ]; + + meta = { + description = "Tiny PyTorch library for maintaining a moving average of a collection of parameters"; + homepage = "https://github.com/fadel/pytorch_ema"; + changelog = "https://github.com/fadel/pytorch_ema/blob/${finalAttrs.src.tag}/CHANGELOG.md"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 66728b94791e..d0f5a23b99e1 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19913,6 +19913,8 @@ self: super: with self; { torch-einops-utils = callPackage ../development/python-modules/torch-einops-utils { }; + torch-ema = callPackage ../development/python-modules/torch-ema { }; + torch-geometric = callPackage ../development/python-modules/torch-geometric { }; torch-memory-saver = callPackage ../development/python-modules/torch-memory-saver { };