fix a bug in eval_batch_retrieval (#9089)

2020-12-15 05:46:55 -08:00
parent c19d04623e
commit 44c340f45f
1 changed files with 1 additions and 1 deletions
--- 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,