izer_max_length,
)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
    num_workers=args.dataloader_num_workers,
) Lastly, the training loop takes care of the remaining steps such as converting images to latent space, adding noise to the input, predicting the noise residual, and calculati