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

            * int(config["gradient_accumulation_steps"])
    )
    print("***** Running training *****")
    print(f"  Num examples = {len(train_dataset)}")
    print(f"  Num Epochs = {num_train_epochs}")
