Apply transforms during __getitem__ in CachedDataset

src/ptbench/data/dataset.py

@@ -333,12 +333,12 @@ class CachedDataset(torch.utils.data.Dataset):
 
         logger.info(f"Caching {self.subset} samples")
         for sample in tqdm(self._samples):
-            sample["data"] = self.transforms(
-                self.raw_data_loader(sample["data"])
-            )
+            sample["data"] = self.raw_data_loader(sample["data"])
 
     def __getitem__(self, idx):
-        return self._samples[idx]
+        sample = self._samples[idx].copy()
+        sample["data"] = self.transforms(sample["data"])
+        return sample
 
     def __len__(self):
         return len(self._samples)