def calculate_perceptual_loss_on_pretrained(
    model: nn.Module,
    state: Any,
    real_images: jnp.ndarray,
    fake_images: jnp.ndarray,
    perceptual_loss_on_logit: bool = False):
  """Calculates perceptual loss on pre-trained model."""
  real_pools, real_outputs = pretrained_model_utils.get_pretrained_embs(
      state, model, images=real_images)
  fake_pools, fake_outputs = pretrained_model_utils.get_pretrained_embs(
      state, model, images=fake_images)
  if perceptual_loss_on_logit:
    loss = l2_loss(real_outputs, fake_outputs)
  else:
    loss = l2_loss(real_pools, fake_pools)
  return loss
Calculates perceptual loss on pre-trained model.
videogvt/train_lib/losses.py