Update chatglm_llm.py (#316)

This commit is contained in:
RainGather 2023-05-11 14:36:52 +08:00 committed by GitHub
parent ecc0c43ab9
commit b0714eed5e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions

View File

@ -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: