def run_frechet_distance(params,
                         batch_iter_1,
                         batch_iter_2,
                         *,
                         num_samples=10000,
                         num_repeats=4,
                         **model_args):
  """Run the I3D model with pmap on two batch iterators and get distance."""
  embeddings_1 = run_frechet_embeddings(
      params, batch_iter_1, num_samples=num_samples * num_repeats,
      **model_args)
  embeddings_2 = run_frechet_embeddings(
      params, batch_iter_2, num_samples=num_samples * num_repeats,
      **model_args)
  scores = frechet_distance_from_logits(embeddings_1['logits_mean'],
                                        embeddings_2['logits_mean'],
                                        num_repeats)
  return scores
Run the I3D model with pmap on two batch iterators and get distance.
videogvt/train_lib/frechet_distance.py