From bfb3ea717917e39cc7de6e2e9c49e2a952d1b623 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Wed, 7 Feb 2024 02:57:54 +1100 Subject: [PATCH] Update save.py --- unsloth/save.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/unsloth/save.py b/unsloth/save.py index 60986a4af8..80219b1851 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -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.")