def strip_train_state_4eval(train_state: TrainState) -> TrainState:
  """Strips the train_state to bare minimum for evaluation."""
  train_state_type = type(train_state)
  if train_state_type == train_state_manager.VQGANTrainState:
    deny_list = ['g_params', 'g_model_state', 'ema_params']
  elif train_state_type == train_state_manager.ScenicTrainState:
    deny_list = ['params', 'model_state']
  else:
    raise ValueError('train_state is not supported.')
  return _strip_train_state_4eval(train_state, deny_list)
Strips the train_state to bare minimum for evaluation.
videogvt/train_lib/eval_utils.py