        progress_bar,
        max_train_steps,
        save_dir
    model.train()
    total_loss = 0

    for step, batch in enumerate(dataloader):
        with accelerator.accumulate(model):
            # Forward pass
            outputs = model(**batch)
            loss = outputs.loss
