From f93fd8e0428c3941782cc64cd36d9e521f7f753e Mon Sep 17 00:00:00 2001
From: Cornelius <39997278+cornzz@users.noreply.github.com>
Date: Wed, 21 Aug 2024 18:31:19 +0200
Subject: [PATCH] Fix generation not concluding when eos was hit for each
 prompt

---
 src/mistral_inference/generate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mistral_inference/generate.py b/src/mistral_inference/generate.py
index c9e35c5..4d2dbe8 100644
--- a/src/mistral_inference/generate.py
+++ b/src/mistral_inference/generate.py
@@ -112,7 +112,7 @@ def generate(
         next_token = sample(last_token_prelogits, temperature=temperature, top_p=0.8)
 
         if eos_id is not None:
-            is_finished = is_finished ^ (next_token == eos_id).cpu()
+            is_finished = is_finished | (next_token == eos_id).cpu()
 
         if is_finished.all():
             break