diff --git a/pts/dataset/loader.py b/pts/dataset/loader.py
index 4a21409..d74cedd 100644
--- a/pts/dataset/loader.py
+++ b/pts/dataset/loader.py
@@ -1,5 +1,5 @@
 from typing import Callable, Iterable, Iterator, List, Optional
-from torch.utils import data
+import itertools
 
 from torch.utils.data import IterableDataset