small bug fix to make our dataset classes subclass torch dataset (#122)

src/utils/data_utils.py

@@ -9,7 +9,7 @@
 from utils.corruptions import corrupt_mapping
 
 
-class CacheDataset:
+class CacheDataset(Dataset):
     """
     Caches the entire dataset in memory.
     """
@@ -27,7 +27,7 @@ def __len__(self):
         return len(self.data)
 
 
-class TransformDataset:
+class TransformDataset(Dataset):
     """
     Applies a transformation to the dataset.
     """
@@ -45,7 +45,7 @@ def __len__(self):
         return len(self.dset)
 
 # Custom dataset wrapper to apply corruption
-class CorruptDataset:
+class CorruptDataset(Dataset):
     def __init__(self, dset: CacheDataset, corruption_fn_name, severity: int = 1):
         print("Initialized CorruptDataset with corruption_fn_name: ", corruption_fn_name)
         self.dset = dset # Original dataset