From 40dd31752485cd643d8c37d0e5826c8557cbd4bf Mon Sep 17 00:00:00 2001 From: Lorenzo Stella Date: Tue, 22 Aug 2023 16:25:30 +0200 Subject: [PATCH] fix ArrowDecoder.decode to return instead of yield --- src/gluonts/dataset/arrow/dec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gluonts/dataset/arrow/dec.py b/src/gluonts/dataset/arrow/dec.py index 3082b70790..148d5311c7 100644 --- a/src/gluonts/dataset/arrow/dec.py +++ b/src/gluonts/dataset/arrow/dec.py @@ -32,7 +32,7 @@ def from_schema(cls, schema): ) def decode(self, batch, row_number: int): - yield from self.decode_batch(batch.slice(row_number, row_number + 1)) + return next(self.decode_batch(batch.slice(row_number, row_number + 1))) def decode_batch(self, batch): for row in batch.to_pandas().to_dict("records"):