From 0eee50537115e12765f9ea71e5a81c02b42a2005 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 3 Mar 2025 17:12:56 -0800 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index cca77bb60b..0f0d4c159f 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -755,7 +755,8 @@ def offload_to_disk(W, model, name, temporary_location : str = "_unsloth_tempora filename = os.path.join(file_location, f"{name}.pt") W = W.weight if hasattr(W, "weight") else W torch.save(W, filename, pickle_module = pickle, pickle_protocol = pickle.HIGHEST_PROTOCOL,) - offloaded_W = torch.load(filename, map_location = "cpu", mmap = True) + # We must use weights_only = False due to pickling + offloaded_W = torch.load(filename, map_location = "cpu", mmap = True, weights_only = False) offloaded_W._offloaded_file_location = filename return offloaded_W pass