fix missing return dict (#8653)

huggingface · Nov 19, 2020 · 62cd9ce · 62cd9ce
1 parent 0c2677f
commit 62cd9ce
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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()}