def get_DIB_data(Dataset):
    """Function that modifies a pytorch dataset so that the target contains the example index."""

    class DIBData(Dataset):
        def __getitem__(self, index):
            img, target = super().__getitem__(index)
            return img, (target, index)  # append the index

    return DIBData