def create_pids2idxs(dataset):
    """Creates a mapping between pids and indexes of images for that pid.
    Returns:
        2D List with pids => idx
    """
    pid2idxs = {}
    for idx, data in enumerate(dataset.data):
        pid = data['pid']
        if pid not in pid2idxs:
            pid2idxs[pid] = [idx]
        else:
            pid2idxs[pid].append(idx)
    return pid2idxs