    )
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / int(config["gradient_accumulation_steps"])
    )
    max_train_steps = config["num_train_epochs"] * num_update_steps_per_epoch

    accelerator.init_trackers(
        config["wandb_project_name"], {"lr_scheduler_type": config["lr_scheduler_type"]}
    )
    total_batch_size = (
            config["train_batch_size"]
