Skip to content

Commit 6dff127

Browse files
pcuencaPrathik Rao
authored andcommitted
Fix training push_to_hub (unconditional image generation): models were not saved before pushing to hub (huggingface#868)
Fix: models were not saved before pushing to hub.
1 parent 6d89d09 commit 6dff127

File tree

1 file changed

+1
-2
lines changed

1 file changed

+1
-2
lines changed

examples/unconditional_image_generation/train_unconditional.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,10 +189,9 @@ def transforms(examples):
189189

190190
if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
191191
# save the model
192+
pipeline.save_pretrained(args.output_dir)
192193
if args.push_to_hub:
193194
repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=False)
194-
else:
195-
pipeline.save_pretrained(args.output_dir)
196195
accelerator.wait_for_everyone()
197196

198197
accelerator.end_training()

0 commit comments

Comments
 (0)