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

Update save.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-07 02:57:54 +11:00
commit bfb3ea7179
1 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

10
unsloth/save.py
View file

@ -377,11 +377,11 @@ def unsloth_save_model(
# Save to GPU memory
state_dict[name] = W
# [TODO] Saving to RAM seems to leak memory???
# elif (max_ram - W.nbytes) > 0:
# # Save to CPU memory
# logger.warning_once(f"We will save to RAM and not VRAM now.")
# state_dict[name] = W.to("cpu", non_blocking = True)
# max_ram = max(max_ram - W.nbytes, 0)
elif (max_ram - W.nbytes) > 0:
# Save to CPU memory
logger.warning_once(f"We will save to RAM and not VRAM now.")
state_dict[name] = W.to("cpu", non_blocking = True, copy = True)
max_ram = max(max_ram - W.nbytes, 0)
else:
# Save to Disk
logger.warning_once(f"We will save to Disk and not RAM now.")