diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py
index aab5dd45..018e17a4 100644
--- a/python/infinilm/base_config.py
+++ b/python/infinilm/base_config.py
@@ -1,7 +1,5 @@
 import argparse
 import json
-import os
-import sys
 import warnings
 
 
@@ -277,7 +275,7 @@ def get_device_str(self, device):
             "nvidia": "cuda",
             "qy": "cuda",
             "cambricon": "mlu",
-            "ascend": "ascend",
+            "ascend": "npu",
             "metax": "cuda",
             "moore": "musa",
             "iluvatar": "cuda",
diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py
index 953b35e5..0149d715 100644
--- a/python/infinilm/llm/llm.py
+++ b/python/infinilm/llm/llm.py
@@ -142,7 +142,7 @@ def __init__(self, config: EngineConfig):
 
     def _init_device(self):
         """Initialize infinicore device and dtype."""
-        supported_devices = ["cpu", "cuda", "mlu", "musa"]
+        supported_devices = ["cpu", "cuda", "mlu", "musa", "npu"]
         device_str = self.config.device
         if device_str not in supported_devices:
             raise ValueError(
@@ -709,13 +709,13 @@ def add_request(
         elif prompt is not None:
             prompt_token_ids = self.engine.tokenize(prompt)
         else:
-            assert messages is not None, (
-                "Either messages or prompt/prompt_token_ids must be provided"
-            )
+            assert (
+                messages is not None
+            ), "Either messages or prompt/prompt_token_ids must be provided"
 
-            assert apply_chat_template, (
-                "apply_chat_template needs to be true for multi-role conversation"
-            )
+            assert (
+                apply_chat_template
+            ), "apply_chat_template needs to be true for multi-role conversation"
 
             prompt = self.engine.apply_chat_template(
                 messages, add_generation_prompt=add_generation_prompt