mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Fix bug in --cpu-memory

Browse source
This commit is contained in:
oobabooga 2023-05-12 06:17:07 -03:00
parent fcb46282c5
commit 2eeb27659d
2 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

3
modules/GPTQ_loader.py
View file

@ -185,7 +185,8 @@ def load_quantized(model_name):
max_memory = {}
for i in range(len(memory_map)):
max_memory[i] = f'{memory_map[i]}GiB' if not re.match('.*ib$', memory_map[i].lower()) else memory_map[i]
max_memory['cpu'] = max_cpu_memory
max_memory['cpu'] = f'{max_cpu_memory}GiB' if not re.match('.*ib$', max_cpu_memory.lower()) else max_cpu_memory
else:
max_memory = accelerate.utils.get_balanced_memory(model)