Fix sampling not ending on EOS token.

PiperOrigin-RevId: 616827061 Change-Id: Ie14c6954669e7c71120c33f4f7eaf7d1d006906d
jonmarrs · Mar 18, 2024 · 4ecc4e7 · 4ecc4e7
1 parent 9937890
commit 4ecc4e7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gemma/sampler.py b/gemma/sampler.py
@@ -154,7 +154,7 @@ def _sample_step(
       logits_buffer = sampler_state.logits_buffer
 
     done = sampler_state.done | jnp.equal(
-        sampler_state.token_buffer[:, decoding_step + 1], self.vocab.eos_id()
+        token_buffer[:, decoding_step + 1], self.vocab.eos_id()
     )
 
     return _SamplingState(