From 98201c223b31cc35e6f91e1b13f2618f0453bed0 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Fri, 18 Dec 2020 13:25:24 +0100
Subject: [PATCH] use class

---
 pts/dataset/loader.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

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,
             )