def load_state_dict(self, state_dict: dict):
    if ('local_epoch' in state_dict['state']):
        self.state_averager.local_epoch = state_dict['state'].pop('local_epoch')
    return self.state_averager.optimizer.load_state_dict(state_dict)
def load_state_dict(self, state_dict: dict):
    if ('local_epoch' in state_dict['state']):
        self.state_averager.local_epoch = state_dict['state'].pop('local_epoch')
    return self.state_averager.optimizer.load_state_dict(state_dict)