            "weight_decay": 0.0,
        },
    ]
    optimizer = torch.optim.AdamW(
        optimizer_grouped_parameters, lr=float(config["learning_rate"]),  # fused=True

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