Fix batched output in decoder-only models

jncraton · Jan 3, 2024 · 3eecebb · 3eecebb
1 parent ab7959b
commit 3eecebb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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]