From f1be0c29d44560195bec9a669cfc04902f7add4b Mon Sep 17 00:00:00 2001 From: Tom Dong Date: Thu, 22 Jul 2021 17:12:41 -0700 Subject: [PATCH] Move input to biencoder to the same device --- blink/main_dense.py | 1 + 1 file changed, 1 insertion(+) diff --git a/blink/main_dense.py b/blink/main_dense.py index c9c01469..557b1bdc 100644 --- a/blink/main_dense.py +++ b/blink/main_dense.py @@ -241,6 +241,7 @@ def _run_biencoder(biencoder, dataloader, candidate_encoding, top_k=100, indexer all_scores = [] for batch in tqdm(dataloader): context_input, _, label_ids = batch + context_input = context_input.to(device=biencoder.device) with torch.no_grad(): if indexer is not None: context_encoding = biencoder.encode_context(context_input).numpy()