fix device management

deepset-ai · Aug 28, 2023 · 72a9a53 · 72a9a53
1 parent fa0a8c6
commit 72a9a53
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/haystack/preview/components/embedders/sentence_transformers_document_embedder.py b/haystack/preview/components/embedders/sentence_transformers_document_embedder.py
@@ -41,7 +41,7 @@ def __init__(
 
         self.model_name_or_path = model_name_or_path
         # TODO: remove device parameter and use Haystack's device management once migrated
-        self.device = device
+        self.device = device or "cpu"
         self.use_auth_token = use_auth_token
         self.batch_size = batch_size
         self.progress_bar = progress_bar
@@ -56,7 +56,7 @@ def to_dict(self) -> Dict[str, Any]:
         return default_to_dict(
             self,
             model_name_or_path=self.model_name_or_path,
-            device=str(self.device) if self.device else None,
+            device=self.device,
             use_auth_token=self.use_auth_token,
             batch_size=self.batch_size,
             progress_bar=self.progress_bar,