Unverified Commit 62cd9ce9 authored by Quentin Lhoest's avatar Quentin Lhoest Committed by GitHub
Browse files

fix missing return dict (#8653)

parent 0c2677f5
...@@ -47,7 +47,7 @@ def embed(documents: dict, ctx_encoder: DPRContextEncoder, ctx_tokenizer: DPRCon ...@@ -47,7 +47,7 @@ def embed(documents: dict, ctx_encoder: DPRContextEncoder, ctx_tokenizer: DPRCon
input_ids = ctx_tokenizer( input_ids = ctx_tokenizer(
documents["title"], documents["text"], truncation=True, padding="longest", return_tensors="pt" documents["title"], documents["text"], truncation=True, padding="longest", return_tensors="pt"
)["input_ids"] )["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()} return {"embeddings": embeddings.detach().cpu().numpy()}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment