Fix bug in output_iterator when prediction is terminal

Author: aron

replicate/prediction.py

@@ -248,6 +248,11 @@ def output_iterator(self) -> Iterator[Any]:
         """
         Return an iterator of the prediction output.
         """
+        if (
+            self.status in ["succeeded", "failed", "canceled"]
+            and self.output is not None
+        ):
+            yield from self.output
 
         # TODO: check output is list
         previous_output = self.output or []
@@ -270,6 +275,12 @@ async def async_output_iterator(self) -> AsyncIterator[Any]:
         """
         Return an asynchronous iterator of the prediction output.
         """
+        if (
+            self.status in ["succeeded", "failed", "canceled"]
+            and self.output is not None
+        ):
+            for item in self.output:
+                yield item
 
         # TODO: check output is list
         previous_output = self.output or []