zhengjl/text-generation-webui
Public Access
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Use the CPU if no GPU is detected

Browse Source
This commit is contained in:
oobabooga
2023-02-11 23:17:06 -03:00
parent 337290777b
commit 5d3f15b915
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
server.py
View File

@@ -120,6 +120,9 @@ def load_model(model_name):
else: else:
command = "AutoModelForCausalLM.from_pretrained" command = "AutoModelForCausalLM.from_pretrained"
params = ["low_cpu_mem_usage=True"] params = ["low_cpu_mem_usage=True"]
if not args.cpu and not torch.cuda.is_available():
print("Warning: no GPU has been detected.\nFalling back to CPU mode.\n")
args.cpu = True
if args.cpu: if args.cpu:
params.append("low_cpu_mem_usage=True") params.append("low_cpu_mem_usage=True")