From 20bc7d4e04650ee7de91654f2e6af5b63c7c5b62 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Sat, 19 Dec 2020 12:09:22 +0100
Subject: [PATCH] return iter

---
 pts/dataset/loader.py  | 5 +++--
 pts/model/estimator.py | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/pts/dataset/loader.py b/pts/dataset/loader.py
index 2fb2c1d..92fd76d 100644
--- a/pts/dataset/loader.py
+++ b/pts/dataset/loader.py
@@ -7,6 +7,7 @@ from gluonts.dataset.common import Dataset
 from gluonts.transform import Transformation, TransformedDataset
 from gluonts.itertools import Cyclic, PseudoShuffled
 
+
 class TransformedIterableDataset(IterableDataset):
     def __init__(
         self,
@@ -26,10 +27,10 @@ class TransformedIterableDataset(IterableDataset):
 
     def __iter__(self):
         if self.shuffle_buffer_length is None:
-            return self.transformed_dataset.__iter__()
+            return iter(self.transformed_dataset)
         else:
             shuffled = PseudoShuffled(
                 self.transformed_dataset,
                 shuffle_buffer_length=self.shuffle_buffer_length,
             )
-            return shuffled.__iter__()
+            return iter(shuffled)
diff --git a/pts/model/estimator.py b/pts/model/estimator.py
index 9ac8bd7..5ba9b7b 100644
--- a/pts/model/estimator.py
+++ b/pts/model/estimator.py
@@ -131,7 +131,7 @@ class PyTorchEstimator(Estimator):
             trained_net=trained_net,
             predictor=self.create_predictor(
                 transformation, trained_net, self.trainer.device
-            )
+            ),
         )
 
     def train(