mirror of
https://github.com/huggingface/transformers.git
synced 2025-08-03 03:31:05 +06:00
fix missing return dict (#8653)
This commit is contained in:
parent
0c2677f529
commit
62cd9ce9f8
@ -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()}
|
||||||
|
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user