瀏覽代碼

add header

Abhilash Majumder 1 年之前
父節點
當前提交
95e845cb24
共有 1 個文件被更改,包括 1 次插入0 次删除
  1. 1 0
      inference/chat_completion.py

+ 1 - 0
inference/chat_completion.py

@@ -14,6 +14,7 @@ from transformers import LlamaConfig, LlamaTokenizer, LlamaForCausalLM
 from safety_utils import get_safety_checker
 from safety_utils import get_safety_checker
 from model_utils import load_model, load_peft_model
 from model_utils import load_model, load_peft_model
 from chat_utils import read_dialogs_from_file, format_tokens
 from chat_utils import read_dialogs_from_file, format_tokens
+from accelerate.utils import is_xpu_available
 
 
 def main(
 def main(
     model_name,
     model_name,