python3Packages.torch-ema: init at 0.3.0

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 ];
+  };
+})

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 { };