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.
This commit is contained in:
parent
52e8fdb8ae
commit
4dce37432b
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue