def get_pretrained_embs(state: ModelState, model: nn.Module,
                        images: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
  """Extract embeddings from a pretrained model.

  Args:
    state: ModelState containing model parameters.
    model: Pretrained Flax model.
    images: Array of shape (H, W, 3).

  Returns:
    pool: Pooled outputs from intermediate layer of shape (H', W', C).
    outputs: Outputs from last layer with shape (num_classes,).
  """

  if len(images.shape) != 4 or images.shape[3] != 3:
    raise ValueError("images should be of shape (H, W, 3).")
  if images.shape[1] != RESNET_IMG_SIZE and images.shape[2] != RESNET_IMG_SIZE:
    images = jax.image.resize(
        images,
        (images.shape[0], RESNET_IMG_SIZE, RESNET_IMG_SIZE, images.shape[3]),
        "bilinear")
  variables = {"params": state.params, "batch_stats": state.batch_stats}
  pool, outputs = model.apply(variables, images, mutable=False, train=False)
  return pool, outputs
Extract embeddings from a pretrained model.

  Args:
    state: ModelState containing model parameters.
    model: Pretrained Flax model.
    images: Array of shape (H, W, 3).

  Returns:
    pool: Pooled outputs from intermediate layer of shape (H', W', C).
    outputs: Outputs from last layer with shape (num_classes,).
  
videogvt/train_lib/pretrained_model_utils.py