updated mnist

Author: Johanmkr

CollaborativeCoding/dataloaders/mnist_4_9.py

@@ -43,7 +43,12 @@ def __init__(
         self.labels_path = self.mnist_path / (
             MNIST_SOURCE["train_labels"][1] if train else MNIST_SOURCE["test_labels"][1]
         )
-
+        
+        # Functions to map the labels from (4,9) -> (0,5) for CrossEntropyLoss to work properly. 
+        self.label_shift = lambda x: x-4
+        self.label_restore = lambda x: x+4
+        
+        
     def __len__(self):
         return len(self.samples)
 
@@ -66,4 +71,4 @@ def __getitem__(self, idx):
         if self.transform:
             image = self.transform(image)
 
-        return image, label
+        return image, self.label_shift(label)