Onehot encode labels in dataset

Author: c-salomonsen

utils/dataloaders/usps_0_6.py

@@ -106,6 +106,12 @@ def __getitem__(self, idx):
 
         data = data.reshape(16, 16)
 
+        # one hot encode the target
+        target = np.eye(self.num_classes, dtype=np.float32)[target]
+
+        # Add channel dimension
+        data = np.expand_dims(data, axis=0)
+
         if self.transform:
             data = self.transform(data)