From 2cfc3971b63998a42b7d354b40a2293207e3d81d Mon Sep 17 00:00:00 2001 From: Dominic Baggott Date: Mon, 22 May 2023 20:31:51 +0100 Subject: [PATCH] Fix Prediction.output_iterator The `output` variable is shadowed by the yielding loop, so when we hit the line: previous_output = output the value of `output` refers to the last value of `output` from the loop above, not the total output we received. That causes the client to shrink the output down to a single string (which, inconveniently, also is a valid argument to `len()`), which means we take the wrong slice of the output the next time through the loop. Signed-off-by: Dominic Baggott --- replicate/prediction.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/replicate/prediction.py b/replicate/prediction.py index a69217d8..dd7a593c 100644 --- a/replicate/prediction.py +++ b/replicate/prediction.py @@ -33,8 +33,7 @@ def output_iterator(self) -> Iterator[Any]: while self.status not in ["succeeded", "failed", "canceled"]: output = self.output or [] new_output = output[len(previous_output) :] - for output in new_output: - yield output + yield from new_output previous_output = output time.sleep(self._client.poll_interval) self.reload()