Hi, unfortunately, the solution is not efficient since it searches for positive/negative pairs randomly. This takes a long time when the dataset is large. To improve the speed for finding pairs, you can use a dictionary storing the correspondences.

class ContrastiveDataset(Dataset):

    def __init__(self, dataset, positive_prob=0.5):
        super().__init__()
        self.dataset = dataset
        self.positive_prob = positive_prob

        self.h = {}
        # construct the hash table for the correspondenses
        for i, im in enumerate(self.dataset.imgs):
            # im[0] is the image address, and im[1] is the label
            lbl = im[1]
            if lbl in self.h:
                self.h[lbl].append(i)
            else:
                self.h[lbl] = [i]

    def __getitem__(self, index):
        same_class = random.uniform(0, 1)
        same_class = same_class > self.positive_prob
        img_0, label_0 = self.dataset[index]

        class_samples = self.h[label_0]
        if same_class:
            while True:
                rnd_idx = random.randint(0, len(class_samples) - 1)
                index_1 = class_samples[rnd_idx]
                if index_1 != index:
                    img_1, label_1 = self.dataset[index_1]
                    break
        else:
            while True:
                index_1 = random.randint(0, self.__len__() - 1)
                if index_1 != index:
                    img_1, label_1 = self.dataset[index_1]
                    break

        return (img_0, label_0), (img_1, label_1)

    def __len__(self):
        return len(self.dataset)