| from PIL import Image | |
| from torch.utils.data import Dataset | |
| # We have to make a custom dataset class to load them with the torch DataLoader | |
| # Custom dataset class for CIFAR-10 images | |
| class CustomCIFAR10Dataset(Dataset): | |
| def __init__(self, images, labels, transform=None): | |
| self.images = images | |
| self.labels = labels | |
| self.transform = transform | |
| def __len__(self): | |
| return len(self.images) | |
| def __getitem__(self, index): | |
| image = self.images[index] | |
| label = self.labels[index] | |
| # Apply the transformations (if any) | |
| if self.transform is not None: | |
| image = self.transform(image) | |
| return image, label | |