fix a bug in eval_batch_retrieval (#9089)

44c340f4 · Yoshitomo Matsubara · GitHub · c19d0462 · 44c340f4
Unverified Commit 44c340f4 authored Dec 15, 2020 by Yoshitomo Matsubara Committed by GitHub Dec 15, 2020
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

examples/research_projects/rag/eval_rag.py examples/research_projects/rag/eval_rag.py +1 -1

No files found.
--- a/examples/research_projects/rag/eval_rag.py
+++ b/examples/research_projects/rag/eval_rag.py
@@ -96,7 +96,7 @@ def evaluate_batch_retrieval(args, rag_model, questions):
    )["input_ids"].to(args.device)

    question_enc_outputs = rag_model.rag.question_encoder(retriever_input_ids)
-    question_enc_pool_output = question_enc_outputs.pooler_output
+    question_enc_pool_output = question_enc_outputs[0]

    result = rag_model.retriever(
        retriever_input_ids,