Fix GPU Embed Bug (#19)

Minor fix from bug caused by
#16. We need to place the
tensors on the proper device. Note: add GPU tests in the future, since
all the tests run on CPU at the moment.

lotus/models/e5_model.py

@@ -61,7 +61,7 @@ def embed(self, docs: List[str], **kwargs: Dict[str, Any]) -> np.ndarray:
 
         # Calculating the embedding dimension
         total_docs = len(docs)
-        first_batch = self.tokenizer(docs[:1], return_tensors="pt", padding=True, truncation=True)
+        first_batch = self.tokenizer(docs[:1], return_tensors="pt", padding=True, truncation=True).to(self.device)
         embed_dim = self.model(**first_batch).last_hidden_state.size(-1)
 
         # Pre-allocate a tensor for all embeddings