From 11c76393b7a4cec2f7f535604d558418229b1837 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 1 Jul 2026 15:12:37 +0000
Subject: [PATCH 1/4] python3Packages.torchsnapshot: use finalAttrs

---
 pkgs/development/python-modules/torchsnapshot/default.nix | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pkgs/development/python-modules/torchsnapshot/default.nix b/pkgs/development/python-modules/torchsnapshot/default.nix
index 017494ba6ff5..1d90101eb231 100644
--- a/pkgs/development/python-modules/torchsnapshot/default.nix
+++ b/pkgs/development/python-modules/torchsnapshot/default.nix
@@ -21,7 +21,7 @@
   pytestCheckHook,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "torchsnapshot";
   version = "0.1.0";
   pyproject = true;
@@ -29,7 +29,7 @@ buildPythonPackage rec {
   src = fetchFromGitHub {
     owner = "pytorch";
     repo = "torchsnapshot";
-    tag = version;
+    tag = finalAttrs.version;
     hash = "sha256-F8OaxLH8BL6MPNLFv1hBuVmeEdnEQ5w2Qny6by1wP6k=";
   };
 
@@ -74,7 +74,7 @@ buildPythonPackage rec {
   meta = {
     description = "Performant, memory-efficient checkpointing library for PyTorch applications, designed with large, complex distributed workloads in mind";
     homepage = "https://github.com/pytorch/torchsnapshot/";
-    changelog = "https://github.com/pytorch/torchsnapshot/releases/tag/${version}";
+    changelog = "https://github.com/pytorch/torchsnapshot/releases/tag/${finalAttrs.src.tag}";
     license = lib.licenses.bsd3;
     maintainers = with lib.maintainers; [ GaetanLepage ];
     badPlatforms = [
@@ -82,4 +82,4 @@ buildPythonPackage rec {
       lib.systems.inspect.patterns.isDarwin
     ];
   };
-}
+})

From 59170b8937b60ff2baf1b00208fff3e5e8f4960b Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 1 Jul 2026 15:09:36 +0000
Subject: [PATCH 2/4] python3Packages.torchsnapshot: cleanup, add missing
 dependency

---
 pkgs/development/python-modules/torchsnapshot/default.nix | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pkgs/development/python-modules/torchsnapshot/default.nix b/pkgs/development/python-modules/torchsnapshot/default.nix
index 1d90101eb231..154e10b15dd2 100644
--- a/pkgs/development/python-modules/torchsnapshot/default.nix
+++ b/pkgs/development/python-modules/torchsnapshot/default.nix
@@ -11,6 +11,7 @@
   aiohttp,
   importlib-metadata,
   nest-asyncio,
+  numpy,
   psutil,
   pyyaml,
   torch,
@@ -52,6 +53,7 @@ buildPythonPackage (finalAttrs: {
     aiohttp
     importlib-metadata
     nest-asyncio
+    numpy
     psutil
     pyyaml
     torch

From d454b8b972e51720607c102bd17ff631ae4117f2 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 1 Jul 2026 20:21:13 +0000
Subject: [PATCH 3/4] python3Packages.torchsnapshot: enable __structuredAttrs

---
 pkgs/development/python-modules/torchsnapshot/default.nix | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pkgs/development/python-modules/torchsnapshot/default.nix b/pkgs/development/python-modules/torchsnapshot/default.nix
index 154e10b15dd2..701edf400d0d 100644
--- a/pkgs/development/python-modules/torchsnapshot/default.nix
+++ b/pkgs/development/python-modules/torchsnapshot/default.nix
@@ -26,6 +26,7 @@ buildPythonPackage (finalAttrs: {
   pname = "torchsnapshot";
   version = "0.1.0";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "pytorch";

From 891b3c13a49a57795904cf910ba04a35c314196f Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 1 Jul 2026 20:44:10 +0000
Subject: [PATCH 4/4] python3Packages.torchsnapshot: skip failing test on
 aarch64-linux

---
 pkgs/development/python-modules/torchsnapshot/default.nix | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pkgs/development/python-modules/torchsnapshot/default.nix b/pkgs/development/python-modules/torchsnapshot/default.nix
index 701edf400d0d..abef8ccc4f34 100644
--- a/pkgs/development/python-modules/torchsnapshot/default.nix
+++ b/pkgs/development/python-modules/torchsnapshot/default.nix
@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   buildPythonPackage,
   fetchFromGitHub,
 
@@ -72,6 +73,11 @@ buildPythonPackage (finalAttrs: {
     # torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
     # AssertionError: "Socket Timeout" does not match "wait timeout after 5000ms
     "test_linear_barrier_timeout"
+  ]
+  ++ 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_tensor_copy"
   ];
 
   meta = {