From 4dce37432b8dc6ec88302b6c91b0700338c153e3 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Mon, 17 Oct 2022 15:28:56 +0200 Subject: [PATCH] Fix training push_to_hub (unconditional image generation): models were not saved before pushing to hub (#868) Fix: models were not saved before pushing to hub. --- examples/unconditional_image_generation/train_unconditional.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 8ddbdd77..f9f8c85b 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -184,10 +184,9 @@ def main(args): if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: # save the model + pipeline.save_pretrained(args.output_dir) if args.push_to_hub: repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=False) - else: - pipeline.save_pretrained(args.output_dir) accelerator.wait_for_everyone() accelerator.end_training()