Skip to content

GitLab

"...composable_kernel.git" did not exist on "63f8766206b72ad5a25ce8274343938d4fe35ff9"
Unverified Commit 4dce3743 authored by Pedro Cuenca's avatar Pedro Cuenca Committed by GitHub
Browse files

Fix training push_to_hub (unconditional image generation): models were not... 

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.
parent 52e8fdb8
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 2 deletions
examples/unconditional_image_generation/train_unconditional.py
View file @ 4dce3743
...@@ -184,10 +184,9 @@ def main(args): ...@@ -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()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment