def gather_outputs_with_mask(
    outputs: List[train_utils.Batch],
    *,
    mask_key: str = 'batch_mask',
    num_samples: Optional[int] = None) -> EvalFeatureDictCPU:
  """Gathers outputs using batch_mask and returns a dict."""
  outputs = jax.tree_util.tree_map(lambda *x: np.concatenate(x), *outputs)
  masks = outputs.pop(mask_key) > 0
  outputs = jax.tree_util.tree_map(lambda x: x[masks][:num_samples], outputs)
  return outputs
Gathers outputs using batch_mask and returns a dict.
videogvt/train_lib/metrics_lib.py