@ex.capture()
def get_scheduler(optimizer, t_total, warmup, scheduler_name, grad_acc_steps):
    warmup_steps = int((t_total * warmup))
    scheduler = sched_dict[scheduler_name](optimizer, warmup_steps, t_total)
    scheduler.accumulated = 0
    scheduler.grad_acc_steps = grad_acc_steps
    return scheduler
@ex.capture()
def get_scheduler(optimizer, t_total, warmup, scheduler_name, grad_acc_steps):
    warmup_steps = int((t_total * warmup))
    scheduler = sched_dict[scheduler_name](optimizer, warmup_steps, t_total)
    scheduler.accumulated = 0
    scheduler.grad_acc_steps = grad_acc_steps
    return scheduler