def check_train_state_type(restored_train_state: Mapping[str, Any]):
  """Returns the type of the TrainState."""
  train_state = restored_train_state
  if 'g_optimizer' in train_state and 'd_optimizer' in train_state:
    return VQGANTrainStateDeprecated
  elif 'g_opt_state' in train_state and 'd_opt_state' in train_state:
    return VQGANTrainState
  elif 'opt_state' in train_state and 'model_state' in train_state:
    return ScenicTrainState
  else:
    return 'unknown'
Returns the type of the TrainState.
videogvt/train_lib/train_state_manager.py