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.
2022-10-17 15:28:56 +02:00
parent 52e8fdb8ae
commit 4dce37432b
1 changed files with 1 additions and 2 deletions
@@ -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()