Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-26 01:30:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

Don't use low_cpu_mem_usage and device_map together

Browse Source
This commit is contained in:
oobabooga 2023-01-31 13:24:05 -03:00
parent 824329749d
commit bb77f20a6c
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
server.py
View File

@ -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() model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
# Custom # Custom
else: else:
settings = ["low_cpu_mem_usage=True"]
command = "AutoModelForCausalLM.from_pretrained" command = "AutoModelForCausalLM.from_pretrained"
settings = []
if args.cpu: if args.cpu:
settings.append("low_cpu_mem_usage=True")
settings.append("torch_dtype=torch.float32") settings.append("torch_dtype=torch.float32")
else: else:
settings.append("device_map='auto'") settings.append("device_map='auto'")