have most of the pieces to start training the model and all that’s left is putting everything together.
First, you’ll need an optimizer and a learning rate scheduler:


	Copied
>>> from diffusers.optimization import get_cosine_schedule_with_warmup

>>> optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
>>> lr_scheduler = get_cosine_schedule_with_warmup(
...     optimizer=optimizer,
...     num_warmup_steps=config.lr_warm