Skip to content

Commit

Permalink
Merge pull request #214 from cornzz/main
Browse files Browse the repository at this point in the history
Fix generation not concluding when eos was hit for each prompt
  • Loading branch information
patrickvonplaten authored Aug 22, 2024
2 parents 337e3db + f93fd8e commit 3fd585d
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/mistral_inference/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 3fd585d

Please sign in to comment.