From 19266b3f05b4d2372692e862369b4e9746fe6a4f Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Mon, 7 Jul 2025 09:11:11 +0200
Subject: [PATCH] python3Packages.torchsnapshot: fix torch.load call

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

diff --git a/pkgs/development/python-modules/torchsnapshot/default.nix b/pkgs/development/python-modules/torchsnapshot/default.nix
index 90e4c3f2dbbf..017494ba6ff5 100644
--- a/pkgs/development/python-modules/torchsnapshot/default.nix
+++ b/pkgs/development/python-modules/torchsnapshot/default.nix
@@ -33,6 +33,16 @@ buildPythonPackage rec {
     hash = "sha256-F8OaxLH8BL6MPNLFv1hBuVmeEdnEQ5w2Qny6by1wP6k=";
   };
 
+  # _pickle.UnpicklingError: Weights only load failed.
+  # torchsnapshot needs to adapt to the change of torch.load that occured in 2.6.0:
+  # https://pytorch.org/docs/stable/generated/torch.load.html
+  postPatch = ''
+    substituteInPlace torchsnapshot/io_preparers/object.py \
+      --replace-fail \
+        "torch.load(io.BytesIO(buf))" \
+        "torch.load(io.BytesIO(buf), weights_only=False)"
+  '';
+
   build-system = [
     setuptools
   ];