Update beam_search.py

ef6cb209 · George Sterpu · GitHub · 6a3bcef8 · ef6cb209
Unverified Commit ef6cb209 authored Jan 27, 2020 by George Sterpu Committed by GitHub Jan 27, 2020
Show whitespace changes
Inline Side-by-side

Showing with 7 additions and 3 deletions

official/transformer/v2/beam_search.py official/transformer/v2/beam_search.py +7 -3

No files found.
--- a/official/transformer/v2/beam_search.py
+++ b/official/transformer/v2/beam_search.py
@@ -30,9 +30,13 @@ class SequenceBeamSearchV2(v1.SequenceBeamSearch):
    """Beam search for sequences with highest scores."""
    state, state_shapes = self._create_initial_state(initial_ids, initial_cache)
-    finished_state = tf.while_loop(
+    finished_state = tf.nest.map_structure(
-        self._continue_search, self._search_step, loop_vars=[state],
+      tf.stop_gradient,
-        shape_invariants=[state_shapes], parallel_iterations=1, back_prop=False)
+      tf.while_loop(self._continue_search,
+                    self._search_step,
+                    loop_vars=[state],
+                    shape_invariants=[state_shapes],
+                    parallel_iterations=1))
    finished_state = finished_state[0]
    alive_seq = finished_state[_StateKeys.ALIVE_SEQ]