Update beam_search_v1.py

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(
-        self._continue_search, self._search_step, loop_vars=[state],
-        shape_invariants=[state_shapes], parallel_iterations=1, back_prop=False)
+    finished_state = tf.nest.map_structure(
+      tf.stop_gradient,
+      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]