feat: Device param

amikos-tech · Mar 4, 2024 · 1d0f324 · 1d0f324
1 parent 091e466
commit 1d0f324
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/chromadb/utils/embedding_functions.py b/chromadb/utils/embedding_functions.py
@@ -671,7 +671,10 @@ def __call__(self, input: Documents) -> Embeddings:
 
 class OpenCLIPEmbeddingFunction(EmbeddingFunction[Union[Documents, Images]]):
     def __init__(
-        self, model_name: str = "ViT-B-32", checkpoint: str = "laion2b_s34b_b79k"
+        self,
+        model_name: str = "ViT-B-32",
+        checkpoint: str = "laion2b_s34b_b79k",
+        device: Optional[str] = "cpu",
     ) -> None:
         try:
             import open_clip
@@ -697,6 +700,7 @@ def __init__(
             model_name=model_name, pretrained=checkpoint
         )
         self._model = model
+        self._model.to(device)
         self._preprocess = preprocess
         self._tokenizer = open_clip.get_tokenizer(model_name=model_name)