def METHOD_NAME(self):
    r"""
    This is the same as :class:`torch.optim.Optimizer` :meth:`state_dict`,
    but adds an extra entry to record model averager's step to the checkpoint
    to ensure reload does not cause unnecessary warm up again.
    """
    optim_state_dict = self.optim.METHOD_NAME()
    optim_state_dict["step"] = self.averager.step
    return optim_state_dict