import torch


def create_scheduler(cfg, optimizer, train_size):
    """
    Creates a scheduler with warmup and linear decay.
    """
    total_steps = train_size // \
        cfg.grad_accum_steps * cfg.max_epochs

    warmup_steps = total_steps * cfg.warmup_prop

    def lr_lambda(current_step):
        if current_step < warmup_steps:
            return current_step / max(1, warmup_steps)

        return max(0, total_steps - current_step) / \
            max(1, total_steps - warmup_steps)

    return torch.optim.lr_scheduler.LambdaLR(
        optimizer, lr_lambda)