[WIP] fix greedy

huggingface · Oct 16, 2023 · 1f6cdd0 · 1f6cdd0
1 parent 81f631c
commit 1f6cdd0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/neuron/generation/utils.py b/optimum/neuron/generation/utils.py
@@ -1253,7 +1253,7 @@ def greedy_search(
 
             # pre-process distribution
             # Move to cpu to handle arbitrary logits_processor
-            next_tokens_scores = logits_processor(input_ids.clone().to("cpu")[:, :seq_length], next_token_logits.to("cpu"))
+            next_tokens_scores = logits_processor(input_ids.to("cpu")[:, :seq_length], next_token_logits.to("cpu"))
             next_tokens_scores = next_tokens_scores.to(input_ids.device)
 
             # Store scores, attentions and hidden_states when required