def collate_batch(batch: list) -> dict:
    """Collate function for PyTorch DataLoader class.

    Converts a list of dicts to a dict of lists.
    """
    return dict(zip(batch[0], zip(*[d.values() for d in batch])))