code: def restore_checkpoint(self) -> Tuple[Dict[str, Any], Dict[str, Any]]:
          
          latest_checkpoint = self.find_latest_checkpoint()
  
          if latest_checkpoint is None:
              # No checkpoint to restore, start at 0
              return {}, {}
  
          model_path, training_state_path = latest_checkpoint
  
          # Load the parameters onto CPU, then transfer to GPU.
          # This avoids potential OOM on GPU for large models that
          # load parameters onto GPU then make a new GPU copy into the parameter
          # buffer. The GPU transfer happens implicitly in load_state_dict.
          model_state = torch.load(model_path, map_location=nn_util.device_mapping(-1))
          training_state = torch.load(training_state_path, map_location=nn_util.device_mapping(-1))
          return model_state, training_state
docs: Restores a model from a serialization_dir to the last saved checkpoint.
          This includes a training state (typically consisting of an epoch count and optimizer state),
          which is serialized separately from  model parameters. This function should only be used to
          continue training - if you wish to load a model for inference/load parts of a model into a new
          computation graph, you should use the native Pytorch functions:
          `` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
  
          If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
          this function will do nothing and return empty dicts.
  
          Returns
          -------
          states: Tuple[Dict[str, Any], Dict[str, Any]]
              The model state and the training state.