def save_checkpoint(self, basename):
    """Saves checkpoint (multiple files) to disk for the current training step.

    Saving a checkpoint will overwrite any previous checkpoint saved with the
    same ``basename``. Use differing ``basename`` values to save multiple
    checkpoints or multiple copies of the same checkpoint.

    Args:
      basename: Basename for saving a checkpoint. Full file paths for the saved
          checkpoint will combine the output dir, basename, and relevant file
          extensions (e.g., `.weights.npy.gz`).
    """
    if self._output_dir is None:
      _log('Did not save checkpoint as output_dir is None')
      return

    inputs.save_data_counters(self._output_dir)
    if not self.is_chief:
      _log('Did not save checkpoint as we are not chief.')
      return

    dir_and_basename = os.path.join(self._output_dir, basename)
    pkl_file = dir_and_basename + '.pkl.gz'

    _log('Saving checkpoint to %s' % pkl_file, stdout=False)
    weights = self._model.weights
    if base.N_WEIGHTS_SHARDS > 1:
      weights = self._trainer_per_task[0].accelerated_model_with_loss.weights
      weights = tl.unshard(weights)
    state = self._model.state
    compresslevel = 0 if self._use_memory_efficient_trainer else 2
    # Serialize optimizer slots.
    for i, trainer in enumerate(self._trainer_per_task):
      flat_slots = _flatten_and_remove_empty(trainer.slots)
      tl.np_to_file(self._to_bits(flat_slots),
                    f'{dir_and_basename}.opt_slots{i}.npy.gz',
                    compresslevel=compresslevel)
    # We only need the input signature for the body, not for the loss layers.
    # That part is the same across tasks - take it from the first one.
    input_signature = self._batch_signature[:self._model.n_in]
    flat_weights, flat_state = tl.flatten_weights_and_state(weights, state)
    _, flat_eval_state = tl.flatten_weights_and_state(
        weights, self._eval_model.state)
    tl.np_to_file(self._to_bits(flat_weights),
                  f'{dir_and_basename}.weights.npy.gz',
                  compresslevel=compresslevel)
    d = {
        'step': self.step,
        'flat_weights': compresslevel,  # for compatibility with older format
        'flat_state': flat_state,
        'flat_eval_state': flat_eval_state,
        'history': self._history.to_dict(),
        'slots_per_task': compresslevel,  # for compatibility with older format
        'input_signature': input_signature,
        'version_timestamp': 'Mar-10-2021'  # To update in the future if needed.
    }
    pickle_to_file(d, pkl_file, gzip=True)
    _log('Checkpoint saved in %s' % pkl_file, stdout=False)
Saves checkpoint (multiple files) to disk for the current training step.

    Saving a checkpoint will overwrite any previous checkpoint saved with the
    same ``basename``. Use differing ``basename`` values to save multiple
    checkpoints or multiple copies of the same checkpoint.

    Args:
      basename: Basename for saving a checkpoint. Full file paths for the saved
          checkpoint will combine the output dir, basename, and relevant file
          extensions (e.g., `.weights.npy.gz`).
    
trax/supervised/training.py