Add missing condtion check in tensor creation in greedy search (#1289)

huggingface · Aug 23, 2024 · 1548fe3 · 1548fe3
1 parent dfeb8a2
commit 1548fe3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py
@@ -2195,7 +2195,8 @@ def _sample(
         # keep track of which sequences are already finished
         batch_size, cur_len = input_ids.shape
         this_peer_finished = False
-        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+        if not ignore_eos:
+            unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
         model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
 
         bucket_size = model_kwargs.get("bucket_size", -1)