code: def set_states(self, states):
          
          states = pickle.loads(states)
          if isinstance(states, tuple) and len(states) == 2:
              self.states, self.optimizer = states
          else:
              self.states = states
          self.states_synced = dict.fromkeys(self.states.keys(), False)
docs: Sets updater states.