def METHOD_NAME(self, num_updates):
    """Update the learning rate after each update."""
    if num_updates < self.cfg.warmup_updates:
        self.lr = self.cfg.warmup_init_lr + num_updates * self.lr_step
    else:
        curr_updates = num_updates - self.cfg.warmup_updates
        if self.t_mult != 1:
            i = math.floor(
                math.log(
                    1 - curr_updates / self.period * (1 - self.t_mult), self.t_mult
                )
            )
            t_i = self.t_mult ** i * self.period
            t_curr = (
                curr_updates
                - (1 - self.t_mult ** i) / (1 - self.t_mult) * self.period
            )
        else:
            i = math.floor(curr_updates / self.period)
            t_i = self.period
            t_curr = curr_updates - (self.period * i)
        lr_shrink = self.lr_shrink ** i
        min_lr = self.cfg.min_lr * lr_shrink
        max_lr = self.max_lr * lr_shrink
        self.lr = min_lr + 0.5 * (max_lr - min_lr) * (
            1 + math.cos(math.pi * t_curr / t_i)
        )
    self.optimizer.set_lr(self.lr)
    return self.lr