import ipex in create_shard_model.py (#2470)

intel · Jan 15, 2024 · 79b0e6c · 79b0e6c
1 parent d4a23f6
commit 79b0e6c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/cpu/inference/python/llm/utils/create_shard_model.py b/examples/cpu/inference/python/llm/utils/create_shard_model.py
@@ -2,6 +2,8 @@
 import argparse
 
 from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaTokenizer, T5ForConditionalGeneration
+# Here import ipex for Baichuan loading compatibility, for other models we can ignore this import
+import intel_extension_for_pytorch
 
 # supported models
 MODEL_CLASSES = {