kanshan/diffusers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix training push_to_hub (unconditional image generation): models were not saved before pushing to hub (#868)

Browse Source 
Fix: models were not saved before pushing to hub.
This commit is contained in:
Pedro Cuenca
2022-10-17 15:28:56 +02:00
committed by GitHub
parent 52e8fdb8ae
commit 4dce37432b
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/unconditional_image_generation/train_unconditional.py
+1 -2
View File
@@ -184,10 +184,9 @@ def main(args):
if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
# save the model # save the model
pipeline.save_pretrained(args.output_dir)
if args.push_to_hub: if args.push_to_hub:
repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=False) repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=False)
else:
pipeline.save_pretrained(args.output_dir)
accelerator.wait_for_everyone() accelerator.wait_for_everyone()
accelerator.end_training() accelerator.end_training()