Limit device to be cpu or xpu

intel-staging · May 22, 2024 · 3689a31 · 3689a31
1 parent afcb121
commit 3689a31
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/...ations/embeddings/llama-index-embeddings-ipex-llm/llama_index/embeddings/ipex_llm/base.py b/...ations/embeddings/llama-index-embeddings-ipex-llm/llama_index/embeddings/ipex_llm/base.py
@@ -57,9 +57,15 @@ def __init__(
         callback_manager: Optional[CallbackManager] = None,
         **model_kwargs,
     ):
-        # Make default device be cpu
+        # Set "cpu" as default device
         self._device = device or "cpu"
 
+        if self._device not in ["cpu", "xpu"]:
+            logger.warning(
+                "IpexLLMEmbedding currently only supports device to be 'cpu' or 'xpu',"
+                f"but you have: {self._device}; Use 'cpu' instead."
+            )
+
         cache_folder = cache_folder or get_cache_dir()
 
         if model_name is None: