separated train and test transformation

IdentitySubset.py

@@ -2,12 +2,19 @@
 import torch
 
 class IdentitySubset(torch.utils.data.Dataset):
-    def __init__(self, full_ds, indices, id_mapping):
-        self.full_ds = full_ds
+    def __init__(self, dataset, indices, id_mapping, transform=None):
+        self.dataset = dataset
         self.indices = indices
         self.id_mapping = id_mapping
+        self.transform = transform
+
     def __getitem__(self, idx):
-        img, old_id = self.full_ds[self.indices[idx]]
+        img, old_id = self.dataset[self.indices[idx]]
+
+        if self.transform:
+            img = self.transform(img)
+
         return img, self.id_mapping[old_id.item()]
+    
     def __len__(self):
         return len(self.indices)