def load_checkpoint(self, directory=None, filename=None):
    """Loads model weights and step from a checkpoint on disk.

    Args:
      directory: Directory with the checkpoint (self._output_dir by default).
      filename: Checkpoint file name (model.pkl.gz by default).
    """
    directory = directory or self._output_dir
    if directory is None:
      _log('Not loading as both directory and output_dir are None.',
           stdout=False)
      return
    filename = filename or 'model'
    path = os.path.join(directory, filename)
    pkl_path = path + '.pkl.gz'
    if not tf.io.gfile.exists(pkl_path):
      _log(f'Not loading as checkpoint file does not exist: {pkl_path}',
           stdout=False)
      return
    _log('Loading checkpoint from %s' % pkl_path, stdout=False)
    d = unpickle_from_file(pkl_path, gzip=True)
    # Weights are stored in a separate non-pickled file in the new checkpoint
    # format. We support loading old checkpoints with this hack.
    # TODO(lukaszkaiser): remove the hack when not needed any more.
    if isinstance(d['flat_weights'], int):
      weights = tl.np_from_file(path + '.weights.npy.gz',
                                compresslevel=d['flat_weights'])
      d['flat_weights'] = weights
    else:
      d['flat_weights'] = d['flat_weights']
    # The same holds for optimizer slots.
    if 'slots' in d:  # Old checkpoints had just 'slots' for one task.
      if len(self._tasks) != 1:
        raise ValueError(
            'Can\'t load a single-task checkpoint into a multitask Loop.'
        )
      d['slots_per_task'] = [d['slots']]
    # Read from separate files if optimizer slots are in them.
    if 'slots_per_task' in d and isinstance(d['slots_per_task'], int):
      compresslevel = d['slots_per_task']
      d['slots_per_task'] = []
      for i in range(len(self._trainer_per_task)):
        slots = tl.np_from_file(path + f'.opt_slots{i}.npy.gz',
                                compresslevel=compresslevel)
        d['slots_per_task'].append(slots)
    for (trainer, slots) in zip(self._trainer_per_task, d['slots_per_task']):
      matched_flat_slots = _match_by_shape(
          self._to_bits(_flatten_and_remove_empty(trainer.slots)),
          _flatten_and_remove_empty(slots))
      matched_slots, _ = fastmath.tree_unflatten(
          self._from_bits(matched_flat_slots),
          trainer.slots, copy_from_tree=[None, ()])
      trainer.slots = matched_slots
    self._step = d['step']
    self._history = trax_history.History.from_dict(d['history'])
    # This is self._model.init_from_file but optimized to not re-read.
    input_signature = d['input_signature']
    weights_and_state_sig = self._model.weights_and_state_signature(
        input_signature)
    flat_init_weights, flat_init_state = tl.flatten_weights_and_state(
        self._model.weights, self._model.state)
    if len(d['flat_weights']) < len(flat_init_weights):
      _log('Checkpoint has less weights than the model, loading first ones.')
    matched_weights = _match_by_shape(self._to_bits(flat_init_weights),
                                      d['flat_weights'])
    matched_weights = self._from_bits(matched_weights)
    try:
      restored_state = True
      matched_state = _match_by_shape(self._to_bits(flat_init_state),
                                      d['flat_state'])
      matched_state = self._from_bits(matched_state)
      weights, state = tl.unflatten_weights_and_state(
          matched_weights, matched_state, weights_and_state_sig)
      self._model.state = state
    except IndexError:
      _log('Failed loading model state from checkpoint, loading weights only.')
      restored_state = False
      weights, _ = tl.unflatten_weights_and_state(
          matched_weights, (), weights_and_state_sig, weights_only=True)
    self._model.weights = weights
    self._eval_model.weights = self._model.weights
    # Restore eval model state; note: it's not always the same as train state.
    if restored_state:
      if 'flat_eval_state' in d:
        flat_eval_state = d['flat_eval_state']
      else:  # It wasn't saved in old checkpoints; remove this branch once done.
        flat_eval_state = d['flat_state']
      _, eval_state = tl.unflatten_weights_and_state(
          matched_weights, flat_eval_state, weights_and_state_sig)
      self._eval_model.state = eval_state
    _log('Checkpoint loaded from %s' % pkl_path, stdout=False)
Loads model weights and step from a checkpoint on disk.

    Args:
      directory: Directory with the checkpoint (self._output_dir by default).
      filename: Checkpoint file name (model.pkl.gz by default).
    
trax/supervised/training.py