diff --git a/pts/dataset/loader.py b/pts/dataset/loader.py
index b179238..17e806e 100644
--- a/pts/dataset/loader.py
+++ b/pts/dataset/loader.py
@@ -95,36 +95,6 @@ class DataLoader(Iterable[DataEntry]):
         self.dtype = dtype
 
 
-class TransformedIterableDataset(torch.utils.data.IterableDataset):
-    def __init__(self, dataset, is_train, transform):
-        self.dataset = dataset
-        self.transform = transform
-        self.is_train = is_train
-        self._cur_iter = None
-
-    def _iterate_forever(self, collection: Iterable[DataEntry]) -> Iterator[DataEntry]:
-        # iterate forever over the collection, the collection must be non empty
-        while True:
-            try:
-                first = next(iter(collection))
-            except StopIteration:
-                raise Exception("empty dataset")
-            else:
-                for x in itertools.chain([first], collection):
-                    yield x
-    
-    def __iter__(self):
-        if self._cur_iter is None:
-            self._cur_iter = self.transform(
-                self._iterate_forever(self.dataset), is_train=self.is_train
-            )
-        assert self._cur_iter is not None
-        while True:
-            data_entry = next(self._cur_iter)
-            yield {k:(v.astype(np.float32) if v.dtype.kind == "f" else v) for k,v in data_entry.items() if isinstance(v, np.ndarray)==True}
-
-
-
 class TrainDataLoader(DataLoader):
     """
     An Iterable type for iterating and transforming a dataset, in batches of a
diff --git a/pts/dataset/transformed_iterable_dataset.py b/pts/dataset/transformed_iterable_dataset.py
new file mode 100644
index 0000000..c8f7031
--- /dev/null
+++ b/pts/dataset/transformed_iterable_dataset.py
@@ -0,0 +1,46 @@
+import itertools
+from typing import Dict, Iterable, Iterator, Optional
+
+import numpy as np
+import torch
+
+from pts.transform import Transformation
+
+from .common import DataEntry, Dataset
+
+class TransformedIterableDataset(torch.utils.data.IterableDataset):
+    def __init__(
+        self, dataset: Dataset, is_train: bool, transform: Transformation
+    ) -> None:
+        self.dataset = dataset
+        self.transform = transform
+        self.is_train = is_train
+        self._cur_iter: Optional[Iterator] = None
+
+    def _iterate_forever(self, collection: Iterable[DataEntry]) -> Iterator[DataEntry]:
+        # iterate forever over the collection, the collection must be non empty
+        while True:
+            try:
+                first = next(iter(collection))
+            except StopIteration:
+                raise Exception("empty dataset")
+            else:
+                for x in itertools.chain([first], collection):
+                    yield x
+
+    def __iter__(self) -> Dict[str, np.ndarray]:
+        if self._cur_iter is None:
+            self._cur_iter = self.transform(
+                self._iterate_forever(self.dataset), is_train=self.is_train
+            )
+        assert self._cur_iter is not None
+        while True:
+            data_entry = next(self._cur_iter)
+            yield {
+                k: (v.astype(np.float32) if v.dtype.kind == "f" else v)
+                for k, v in data_entry.items()
+                if isinstance(v, np.ndarray) == True
+            }
+
+    def __len__(self) -> int:
+        return len(self.dataset)
\ No newline at end of file