docs: def _restore_training_state(self, restore_state):
          """Restores the model and optimizer states
  
          This helper function restores the model's state to a given iteration so
          that a user can resume training at any epoch.
  
          Args:
              restore_state: a state_dict dictionary
          """
          self.load_state_dict(restore_state["model"])
          self.optimizer.load_state_dict(restore_state["optimizer"])
          self.lr_scheduler.load_state_dict(restore_state["lr_scheduler"])
          start_iteration = restore_state["iteration"] + 1
          if self.config["verbose"]:
              print(f"Restored checkpoint to iteration {start_iteration}.")
  
          if restore_state["best_model_found"]:
              # Update checkpointer with appropriate information about best model
              # Note that the best model found so far may not be the model in the
              # checkpoint that is currently being loaded.
              self.checkpointer.best_model_found = True
              self.checkpointer.best_iteration = restore_state["best_iteration"]
              self.checkpointer.best_score = restore_state["best_score"]
              if self.config["verbose"]:
                  print(
                      f"Updated checkpointer: "
                      f"best_score={self.checkpointer.best_score:.3f}, "
                      f"best_iteration={self.checkpointer.best_iteration}"
                  )
          return start_iteration