.     image_grid.save(f"{test_dir}/{epoch:04d}.png")
Now you can wrap all these components together in a training loop with 🤗 Accelerate for easy TensorBoard logging, gradient accumulation, and mixed precision training. To upload the model to the Hub, write a function to get your repository name and information and then push it to the Hub.
💡 The training loop below may look intimidating and long, but it’ll be worth it later when you launch your