        range(max_train_steps), disable=not accelerator.is_local_main_process
    )
    completed_steps = 0
    starting_epoch = 0
    for epoch in range(starting_epoch, num_train_epochs):

            model,
            train_dataloader,
            accelerator,
            optimizer,
            lr_scheduler,
