diff --git a/src/ptbench/data/dataset.py b/src/ptbench/data/dataset.py index 243425f1ea9199a152163ad24b711e46cf88b3b6..15dc32a96f6662b5ba8afbcdf06632b142840f79 100644 --- a/src/ptbench/data/dataset.py +++ b/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)