From e154c393fe236068aba6df832f03316e5dc089b2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A1rton=20Kardos?= Date: Tue, 13 Feb 2024 10:19:32 +0100 Subject: [PATCH] Embeddings are sent back to the CPU, so they can be converted to numpy arrays --- src/seb/registered_models/e5_mistral.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/seb/registered_models/e5_mistral.py b/src/seb/registered_models/e5_mistral.py index e66b3c75..c7413809 100644 --- a/src/seb/registered_models/e5_mistral.py +++ b/src/seb/registered_models/e5_mistral.py @@ -183,7 +183,7 @@ def encode( ) batched_embeddings.append(embeddings) - return torch.cat(batched_embeddings) + return torch.cat(batched_embeddings).to("cpu") def encode_corpus(self, corpus: list[dict[str, str]], **kwargs: Any): if isinstance(corpus, dict):