From 3eecebbc7732f74353848d4063a7c7450a6289e2 Mon Sep 17 00:00:00 2001
From: Jon Craton <jncraton@gmail.com>
Date: Wed, 3 Jan 2024 14:48:41 -0500
Subject: [PATCH] Fix batched output in decoder-only models

---
 languagemodels/inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/languagemodels/inference.py b/languagemodels/inference.py
index 9d98b9c..c280147 100644
--- a/languagemodels/inference.py
+++ b/languagemodels/inference.py
@@ -170,7 +170,7 @@ def generate(
             beam_size=1,
             include_prompt_in_result=False,
         )
-        outputs_ids = results[0].sequences_ids[0]
+        outputs_ids = [r.sequences_ids[0] for r in results]
 
     return [tokenizer.decode(i, skip_special_tokens=True).lstrip() for i in outputs_ids]