def step_update(self, num_updates):
        """Update the learning rate after each update."""
        manual_keys = [k for k in self.update2lr if k <= num_updates]
        if manual_keys:
            manual_lr = self.update2lr[max(manual_keys)]
        else:
            logger.warning(
                "epoch={} does not exist in manual lr input update2lr={}...".format(
                    num_updates,
                    list(self.update2lr.items())[
                        : min(10, len(self.update2lr.keys()) - 1)
                    ],
                )
            )
            manual_lr = self.optimizer.get_lr()

        self.optimizer.set_lr(manual_lr)
        return self.optimizer.get_lr()
Update the learning rate after each update.
structured_prompting/fairseq-version/fairseq/fairseq/optim/lr_scheduler/manual_lr_scheduler.py