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
--
GitLab