def METHOD_NAME(self):
    """
    Return a kwarg dictionary that will be used to override optimizer
    args stored in checkpoints. This allows us to load a checkpoint and
    resume training using a different set of optimizer args, e.g., with a
    different learning rate.
    """
    return {
        "lr": self.args.lr[0],
        "betas": eval(self.args.adamax_betas),
        "eps": self.args.adamax_eps,
        "weight_decay": self.args.weight_decay,
        "bias_correction": not self.args.no_bias_correction,
    }