diff --git a/pts/dataset/loader.py b/pts/dataset/loader.py
index d74cedd..7d9cdc5 100644
--- a/pts/dataset/loader.py
+++ b/pts/dataset/loader.py
@@ -5,8 +5,7 @@ from torch.utils.data import IterableDataset
 
 from gluonts.dataset.common import Dataset
 from gluonts.transform import Transformation, TransformedDataset
-from gluonts.itertools import cyclic, pseudo_shuffled
-
+from gluonts.itertools import Cyclic, PseudoShuffled
 
 class TransformedIterableDataset(IterableDataset):
     def __init__(
@@ -20,16 +19,16 @@ class TransformedIterableDataset(IterableDataset):
         self.shuffle_buffer_length = shuffle_buffer_length
 
         self.transformed_dataset = TransformedDataset(
-            cyclic(dataset),
+            Cyclic(dataset),
             transform,
             is_train=is_train,
         )
 
     def __iter__(self):
         if self.shuffle_buffer_length is None:
-            return iter(self.transformed_dataset)
+            return self.transformed_dataset
         else:
-            return pseudo_shuffled(
-                iter(self.transformed_dataset),
+            return PseudoShuffled(
+                self.transformed_dataset,
                 shuffle_buffer_length=self.shuffle_buffer_length,
             )