Update beam_search_v1.py

2e21c80e · George Sterpu · GitHub · ef6cb209 · 2e21c80e
Unverified Commit 2e21c80e 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/nlp/transformer/beam_search_v1.py official/nlp/transformer/beam_search_v1.py +7 -3

No files found.
--- a/official/nlp/transformer/beam_search_v1.py
+++ b/official/nlp/transformer/beam_search_v1.py
@@ -128,9 +128,13 @@ class SequenceBeamSearch(object):
    """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]