Flat - Linear Learning Rate Schedule from Wav2Vec2 paper
Sharing Trainer code for the same learning rate schedule as the paper, 10% warmup, 40% flat, 50% linear decay
def get_flat_linear_schedule_with_warmup(optimizer:Optimizer, num_warmup_steps:int,
                                         num_training_steps:int, last_epoch:int =-1):
    def lr_lambda(current_step):
        constant_steps = int(num_training_steps * 0.4)
        warmup_steps = int(num_training_steps * 0.1)
        if current_step < warmup_steps:
            return float(current_step) / float(max(1, warmup_steps))
        elif current_step < warmup_steps+constant_steps:
            return 1
        else: return max(
            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - (warmup_steps+constant_steps)))
        )
    return LambdaLR(optimizer, lr_lambda, last_epoch)
def get_flat_cheduler(
    name: Union[str, SchedulerType] = None,
    optimizer: Optimizer = None,
    num_warmup_steps: Optional[int] = None,
    num_training_steps: Optional[int] = None,
):
    return get_flat_linear_schedule_with_warmup(optimizer, num_warmup_steps=num_warmup_steps, 
                                                num_training_steps=num_training_steps)

And creating Trainer warpper:
class FlatTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    def create_flat_scheduler(self, num_training_steps: int):
        self.lr_scheduler = get_flat_cheduler(optimizer = self.optimizer,
                                              num_training_steps=num_training_steps)
    def create_optimizer_and_scheduler(self, num_training_steps):
        self.create_optimizer()
        self.create_flat_scheduler(num_training_steps)