def _compute_lpips_one_step(batch_outputs, metric_params):
  """Compute the lpips for one batch of generated videos."""
  # lpips calls tf_hub which cannot be jit/pmaped by tf2jax
  suffix = ''
  cur_batch = batch_outputs
  original_videos = cur_batch[f'original_video{suffix}']
  generated_videos = cur_batch[f'generated_video{suffix}']
  assert generated_videos.ndim == 5  # total_bs, t, h, w, 3
  return eval_utils.run_lpips_models(
      original_videos,
      generated_videos,
      metric_functions=metric_params['lpips'])
Compute the lpips for one batch of generated videos.
videogvt/trainers/maskgvt_trainer.py