From 62cd9ce9f860774e137b7bba5b900b6e68000d9c Mon Sep 17 00:00:00 2001 From: Quentin Lhoest <42851186+lhoestq@users.noreply.github.com> Date: Thu, 19 Nov 2020 15:17:18 +0100 Subject: [PATCH] fix missing return dict (#8653) --- examples/rag/use_own_knowledge_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/rag/use_own_knowledge_dataset.py b/examples/rag/use_own_knowledge_dataset.py index 081216f99f..fd465e6900 100644 --- a/examples/rag/use_own_knowledge_dataset.py +++ b/examples/rag/use_own_knowledge_dataset.py @@ -47,7 +47,7 @@ def embed(documents: dict, ctx_encoder: DPRContextEncoder, ctx_tokenizer: DPRCon input_ids = ctx_tokenizer( documents["title"], documents["text"], truncation=True, padding="longest", return_tensors="pt" )["input_ids"] - embeddings = ctx_encoder(input_ids.to(device=device)).pooler_output + embeddings = ctx_encoder(input_ids.to(device=device), return_dict=True).pooler_output return {"embeddings": embeddings.detach().cpu().numpy()}