return iter

This commit is contained in:
Dr. Kashif Rasul
2020-12-19 12:09:22 +01:00
parent bf0a199191
commit 20bc7d4e04
2 changed files with 4 additions and 3 deletions
+3 -2
View File
@@ -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)
+1 -1
View File
@@ -131,7 +131,7 @@ class PyTorchEstimator(Estimator):
trained_net=trained_net,
predictor=self.create_predictor(
transformation, trained_net, self.trainer.device
)
),
)
def train(