ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update _utils.py

Browse source
This commit is contained in:
Daniel Han 2025-03-03 17:12:56 -08:00
commit 0eee505371
1 changed files with 2 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/models/_utils.py
View file

@ -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