fix: ensure SentenceTransformer models load directly to device

Author: JamieHoldcroft

rankify/retrievers/bge_reasoner_retriever.py

@@ -99,7 +99,7 @@ def __init__(
     def _load_model(self):
         # ReasonEmbed SentenceTransformer models
         if self.model_id == "bge-reasoner-embed":
-            self.model = SentenceTransformer(self.checkpoint or "BAAI/bge-reasoner-embed-qwen3-8b-0923", trust_remote_code=True)
+            self.model = SentenceTransformer(self.checkpoint or "BAAI/bge-reasoner-embed-qwen3-8b-0923", trust_remote_code=True, device=self.device)
             self.model = self.model.to(self.device)
         else:
             raise ValueError(f"The model {self.model_id} is not supported")

rankify/retrievers/diver_dense_retriever.py

@@ -149,18 +149,18 @@ def __init__(
     def _load_model(self):
         # SentenceTransformer models
         if self.model_id == "bge":
-            self.model = SentenceTransformer('BAAI/bge-large-en-v1.5')
+            self.model = SentenceTransformer('BAAI/bge-large-en-v1.5', device=self.device)
         elif self.model_id == "sbert":
-            self.model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
+            self.model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2', device=self.device)
         elif self.model_id == "contriever_st":
-            self.model = SentenceTransformer('nishimoto/contriever-sentencetransformer')
+            self.model = SentenceTransformer('nishimoto/contriever-sentencetransformer', device=self.device)
         elif self.model_id == "nomic":
-            self.model = SentenceTransformer(self.checkpoint or "nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
+            self.model = SentenceTransformer(self.checkpoint or "nomic-ai/nomic-embed-text-v1", trust_remote_code=True, device=self.device)
         elif self.model_id == "inst-l":
-            self.model = SentenceTransformer("hkunlp/instructor-large")
+            self.model = SentenceTransformer("hkunlp/instructor-large", device=self.device)
             self.model.max_seq_length = self.doc_max_length
         elif self.model_id == "inst-xl":
-            self.model = SentenceTransformer("hkunlp/instructor-xl")
+            self.model = SentenceTransformer("hkunlp/instructor-xl", device=self.device)
             self.model.max_seq_length = self.doc_max_length
 
         # HF AutoModel models (sf, e5, rader)

rankify/retrievers/reasonembed_retriever.py

@@ -102,13 +102,13 @@ def __init__(
     def _load_model(self):
         # ReasonEmbed SentenceTransformer models
         if self.model_id == "qwen3-8b":
-            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-qwen3-8b-0928", trust_remote_code=True)
+            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-qwen3-8b-0928", trust_remote_code=True, device=self.device)
             self.model = self.model.to(self.device)
         elif self.model_id == "qwen3-4b":
-            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-qwen3-4b-0928", trust_remote_code=True)
+            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-qwen3-4b-0928", trust_remote_code=True, device=self.device)
             self.model = self.model.to(self.device)
         elif self.model_id == "llama-8b":
-            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-llama-3.1-8b-0928", trust_remote_code=True)
+            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-llama-3.1-8b-0928", trust_remote_code=True, device=self.device)
             self.model = self.model.to(self.device)
         else:
             raise ValueError(f"The model {self.model_id} is not supported")