def METHOD_NAME(self, num_updates):
    """Update the learning rate after each update."""
    stage, steps_in_stage = self._decide_stage(num_updates)
    if stage == 0:
        self.lr = self.init_lr + self.warmup_rate * steps_in_stage
    elif stage == 1:
        self.lr = self.peak_lr
    elif stage == 2:
        self.lr = self.peak_lr * math.exp(-self.decay_factor * steps_in_stage)
    elif stage == 3:
        self.lr = self.final_lr
    else:
        raise ValueError("Undefined stage")
    self.optimizer.set_lr(self.lr)
    return self.lr