Fix beam decoding end-condition bug.

PiperOrigin-RevId: 309795875
google · May 4, 2020 · 16c2e39 · 16c2e39
1 parent d8b7833
commit 16c2e39
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trax/models/beam_search.py b/trax/models/beam_search.py
@@ -251,7 +251,7 @@ def beam_search(batch_size,
   def beam_search_loop_cond_fn(state):
     """Beam search loop termination condition."""
     # Have we reached max decoding length?
-    not_at_end = (state.cur_index <= max_decode_len)
+    not_at_end = (state.cur_index < max_decode_len - 1)
 
     # Is no further progress in the beam search possible?
     # Get the best possible scores from alive sequences.