def get_pretrained_model(
    model_name: str = "resnet50",
    checkpoint_path: Optional[str] = _DEFAULT_RESNET_PATH
) -> Tuple[nn.Module, ModelState]:
  """Returns a pretrained model loaded from weights in checkpoint_dir.

  Args:
    model_name: Name of model architecture to load. Currently only supports
      "resnet50".
    checkpoint_path: Path of .npy containing pretrained state.

  Returns:
    model: Flax nn.Nodule model architecture.
    state: The initialized ModelState with the optimizer.
  """
  if model_name not in VALID_MODELS:
    raise ValueError(f"Model {model_name} not supported.")
  # Initialize model.
  config = ml_collections.ConfigDict()
  config.model_name = "resnet50"
  config.sgd_momentum = 0.9  # Unused for inference.
  config.seed = 42  # Unused for inference.
  model_rng = jax.random.PRNGKey(config.seed)
  model, state = create_train_state(
      config,
      model_rng,
      input_shape=(1, RESNET_IMG_SIZE, RESNET_IMG_SIZE, 3),
      num_classes=1000)

  if checkpoint_path is not None:
    # Set up checkpointing of the model and the input pipeline.
    with tf.io.gfile.GFile(checkpoint_path, "rb") as f:
      checkpoint_data = np.load(f, allow_pickle=True).item()
      state = ModelState(
          params=checkpoint_data["params"],
          batch_stats=checkpoint_data["batch_stats"])
  return model, state
Returns a pretrained model loaded from weights in checkpoint_dir.

  Args:
    model_name: Name of model architecture to load. Currently only supports
      "resnet50".
    checkpoint_path: Path of .npy containing pretrained state.

  Returns:
    model: Flax nn.Nodule model architecture.
    state: The initialized ModelState with the optimizer.
  
videogvt/train_lib/pretrained_model_utils.py