examples/research_projects/bertabs/modeling_bertabs.py/BertSumOptimizer/step
class BertSumOptimizer:    def step(self):
        self._step += 1
        for stack, optimizer in self.optimizers.items():
            new_rate = self._update_rate(stack)
            for param_group in optimizer.param_groups:
                param_group["lr"] = new_rate
            optimizer.step()
            self.current_learning_rates[stack] = new_rate