Merge pull request #8090 from georgesterpu:master

PiperOrigin-RevId: 293174292

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(
-        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]