From 1368f3fca0a4d5c55831325fd652fc650a52ba78 Mon Sep 17 00:00:00 2001
From: JQ <460494839@qq.com>
Date: Thu, 28 Dec 2023 22:38:46 +0800
Subject: [PATCH] NPU needs to be initialized when starting a new process
 (#2843)

---
 fastchat/serve/model_worker.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/fastchat/serve/model_worker.py b/fastchat/serve/model_worker.py
index fa43acf9a..dac6ac8b4 100644
--- a/fastchat/serve/model_worker.py
+++ b/fastchat/serve/model_worker.py
@@ -31,7 +31,6 @@
     str_to_torch_dtype,
 )
 
-
 worker_id = str(uuid.uuid4())[:8]
 logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
 
@@ -103,6 +102,10 @@ def __init__(
             self.init_heart_beat()
 
     def generate_stream_gate(self, params):
+        if self.device == "npu":
+            import torch_npu
+
+            torch_npu.npu.set_device("npu:0")
         self.call_ct += 1
 
         try: