def custom_collate_fn(batch):
    """
    Collate function for DataLoader that collates a list of dicts into a dict of lists.
    """
    collated_batch = {}
    for key in batch[0].keys():
        collated_batch[key] = [item[key] for item in batch]
    return collated_batch
Collate function for DataLoader that collates a list of dicts into a dict of lists.
open_flamingo/eval/utils.py