fix load model hook

huggingface · Jun 17, 2023 · 2a03657 · 2a03657
1 parent 3b7708d
commit 2a03657
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/training/train_muse.py b/training/train_muse.py
@@ -267,7 +267,7 @@ def load_model_hook(models, input_dir):
             load_model = EMAModel.from_pretrained(os.path.join(input_dir, "ema_model"), model_cls=model_cls)
             ema.load_state_dict(load_model.state_dict())
             ema.to(accelerator.device)
-            del ema
+            del load_model
 
         def save_model_hook(models, weights, output_dir):
             ema.save_pretrained(os.path.join(output_dir, "ema_model"))