def apply_batch_mask(batch: Batch) -> Batch:
  """Only keep valid samples in a flattened batch."""
  if 'batch_mask' in batch:
    mask = batch.pop('batch_mask') > 0
    batch = jax.tree_util.tree_map(lambda x: x[mask], batch)
  return batch
Only keep valid samples in a flattened batch.
videogvt/train_lib/eval_utils.py