def batch_index(batch: dict, index):
    """Index into the tensors of a batch.

    Args:
        batch (dict): Batch dictionary.
        index (object): Index.

    Returns:
        dict: `batch` indexed at `index`.
    """
    return {k: batch_index(v, index) for k, v in batch.items()}