Merge branch 'transformers_future' into transformers_4_49

Luca-Calabria · Mar 11, 2025 · 1aa2fcd · 1aa2fcd
2 parents 96c8a32 + 195fdf8
commit 1aa2fcd
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py
@@ -2364,12 +2364,6 @@ def _contrastive_search(
                 )
             # contrastive_search main logic end
 
-            # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
-            model_kwargs = self._update_model_kwargs_for_generation(
-                outputs,
-                model_kwargs,
-                is_encoder_decoder=self.config.is_encoder_decoder,
-            )
             if synced_gpus and this_peer_finished:
                 continue
 
@@ -2387,6 +2381,11 @@ def _contrastive_search(
 
             if streamer is not None:
                 streamer.put(next_tokens.cpu())
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+            )
 
             # increase cur_len
             cur_len = cur_len + 1