use return dict for rag encoder (#9363)

5f7a07c0 · Derrick Blakely · GitHub · ae333d04 · 5f7a07c0
Unverified Commit 5f7a07c0 authored Jan 02, 2021 by Derrick Blakely Committed by GitHub Jan 02, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

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

No files found.
--- a/src/transformers/models/rag/modeling_rag.py
+++ b/src/transformers/models/rag/modeling_rag.py
@@ -1437,7 +1437,7 @@ class RagTokenForGeneration(RagPreTrainedModel):
        batch_size = context_input_ids.shape[0] // n_docs
        encoder = self.rag.generator.get_encoder()
-        encoder_outputs = encoder(input_ids=context_input_ids, attention_mask=context_attention_mask)
+        encoder_outputs = encoder(input_ids=context_input_ids, attention_mask=context_attention_mask, return_dict=True)
        input_ids = torch.full(
            (batch_size * num_beams, 1),