Просмотр исходного кода

Don't use low_cpu_mem_usage and device_map together

oobabooga 3 лет назад
Родитель
Сommit
bb77f20a6c
1 измененных файлов с 3 добавлено и 2 удалено
  1. 3 2
      server.py

+ 3 - 2
server.py

@@ -87,10 +87,11 @@ def load_model(model_name):
             model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
     # Custom
     else:
-        settings = ["low_cpu_mem_usage=True"]
         command = "AutoModelForCausalLM.from_pretrained"
+        settings = []
 
         if args.cpu:
+            settings.append("low_cpu_mem_usage=True")
             settings.append("torch_dtype=torch.float32")
         else:
             settings.append("device_map='auto'")
@@ -374,7 +375,7 @@ if args.chat or args.cai_chat:
             reply = reply[idx + 1 + len(apply_extensions(f"{current}:", "bot_prefix")):]
         else:
             reply = reply[idx + 1 + len(f"{current}:"):]
-        
+
         if check:
             reply = reply.split('\n')[0].strip()
         else: