code: def state_dict(self) -> dict:
          
          hidden_state = {}
  
          if self.optimizer is not None:
              hidden_state['optimizer'] = self.optimizer.state_dict()
  
          for callback in self.callbacks:
              callback.write_state_dict(self.training_info, hidden_state)
  
          return hidden_state
docs: Calculate hidden state dictionary