        optimizer=optimizer,
        num_warmup_steps=int(config["num_warmup_steps"]) * accelerator.num_processes,
        num_training_steps=max_train_steps * accelerator.num_processes,
    )
    if checkpoint_path is not None:

        sceduler_state = torch.load(os.path.join(checkpoint_path, "scheduler.pt"))
        optimizer.load_state_dict(optim_state)
        lr_scheduler.load_state_dict(sceduler_state)
    # model = freeze(model, freeze_other=False, freeze_ff=True, freeze_ff_layers=[31])
    (
