From b0714eed5efb1cf03f024252f0c660bbe53142e1 Mon Sep 17 00:00:00 2001
From: RainGather <3255329+RainGather@users.noreply.github.com>
Date: Thu, 11 May 2023 14:36:52 +0800
Subject: [PATCH] Update chatglm_llm.py (#316)

---
 models/chatglm_llm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/chatglm_llm.py b/models/chatglm_llm.py
index 2290579..49ce97c 100644
--- a/models/chatglm_llm.py
+++ b/models/chatglm_llm.py
@@ -155,7 +155,7 @@ class ChatGLM(LLM):
                     model = PeftModel.from_pretrained(self.model, LLM_LORA_PATH)
                 # 可传入device_map自定义每张卡的部署情况
                 if device_map is None:
-                    device_map = auto_configure_device_map(num_gpus)
+                    device_map = auto_configure_device_map(num_gpus, use_lora)
 
                 self.model = dispatch_model(self.model.half(), device_map=device_map)
         else: