def state_dict(self) -> dict:
    state_dict = self.state_averager.optimizer.state_dict()
    state_dict['state']['local_epoch'] = self.local_epoch
    return state_dict
def state_dict(self) -> dict:
    state_dict = self.state_averager.optimizer.state_dict()
    state_dict['state']['local_epoch'] = self.local_epoch
    return state_dict