Initialize object array correctly

dottxt-ai · Jun 7, 2023 · 4b356b7 · 4b356b7
1 parent 86212b7
commit 4b356b7
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/outlines/models/openai.py b/outlines/models/openai.py
@@ -228,7 +228,11 @@ async def generate_base(prompt: str, samples: int) -> PILImage:
             response = api_response["data"][i]["b64_json"]
             images.append(Image.open(BytesIO(base64.b64decode(response))))
 
-        return np.array(images, dtype="object")
+        array = np.empty((samples,), dtype="object")
+        for idx, image in enumerate(images):
+            array[idx] = image
+
+        return np.atleast_2d(array)
 
     return generate