1 files changed, 1 insertions, 1 deletions

diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 4dd1100..ccec215 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -90,7 +90,7 @@ def lora_strategy_callbacks(
     def on_checkpoint(step, postfix):
         print(f"Saving checkpoint for step {step}...")
 
-        unet_ = accelerator.unwrap_model(unet)
+        unet_ = accelerator.unwrap_model(unet, False)
         unet_.save_attn_procs(checkpoint_output_dir / f"{step}_{postfix}")
         del unet_