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(