Fixed parameter name for logits_processor (#9790)

c37dcff7 · Michael Glass · GitHub · 0d0efd3a · c37dcff7
Unverified Commit c37dcff7 authored Jan 26, 2021 by Michael Glass Committed by GitHub Jan 26, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/transformers/models/rag/modeling_rag.py src/transformers/models/rag/modeling_rag.py +2 -2

No files found.
--- a/src/transformers/models/rag/modeling_rag.py
+++ b/src/transformers/models/rag/modeling_rag.py
@@ -1486,7 +1486,7 @@ class RagTokenForGeneration(RagPreTrainedModel):
                )
            return self.greedy_search(
                input_ids,
-                pre_processor=pre_processor,
+                logits_processor=pre_processor,
                max_length=max_length,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
@@ -1509,7 +1509,7 @@ class RagTokenForGeneration(RagPreTrainedModel):
            return self.beam_search(
                input_ids,
                beam_scorer,
-                pre_processor=pre_processor,
+                logits_processor=pre_processor,
                max_length=max_length,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,