def forward(self, batch: List[tuple]):
        """Overriding the BaseCollateFunction class's forward method because
        for PIRL we need only one augmented batch, as opposed to both, which the
        BaseCollateFunction creates."""
        batch_size = len(batch)

        # list of transformed images
        img_transforms = [
            self.jigsaw(batch[i][0]).unsqueeze_(0) for i in range(batch_size)
        ]
        img = [self.no_augment(batch[i][0]).unsqueeze_(0) for i in range(batch_size)]
        # list of labels
        labels = torch.LongTensor([item[1] for item in batch])
        # list of filenames
        fnames = [item[2] for item in batch]

        # tuple of transforms
        transforms = (torch.cat(img, 0), torch.cat(img_transforms, 0))

        return transforms, labels, fnames
Overriding the BaseCollateFunction class's forward method because
        for PIRL we need only one augmented batch, as opposed to both, which the
        BaseCollateFunction creates.
lightly/data/collate.py