def restore_checkpoint(checkpoint_path: str,
                       train_state: state_mgr.TrainState,
                       step: Optional[int] = None,
                       is_legacy_checkpoint: Optional[bool] = False):
  """Restores a legacy checkpoint.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    step: Step number to load or None to load the latest. If specified,
      checkpoint_path must be a directory.
    is_legacy_checkpoint: whether in the legacy checkpoint format.

  Returns:
    Training state and an int which is the current step.
  """
  if is_legacy_checkpoint:
    train_state = restore_legacy_checkpoint(checkpoint_path, train_state, step)
  else:
    train_state = checkpoints.restore_checkpoint(
        checkpoint_path, train_state, step
    )
  return train_state
Restores a legacy checkpoint.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    step: Step number to load or None to load the latest. If specified,
      checkpoint_path must be a directory.
    is_legacy_checkpoint: whether in the legacy checkpoint format.

  Returns:
    Training state and an int which is the current step.
  
videogvt/train_lib/train_utils.py