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

Use 'torch.backends.mps.is_available' to check if mps is supported (#3164)

Browse Source
This commit is contained in:
appe233 2023-07-18 08:27:18 +08:00 committed by GitHub
parent 234c58ccd1
commit 89e0d15cf5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
modules/LoRA.py
View File

@ -132,7 +132,7 @@ def add_lora_transformers(lora_names):
if not shared.args.load_in_8bit and not shared.args.cpu:
shared.model.half()
if not hasattr(shared.model, "hf_device_map"):
if torch.has_mps:
if torch.backends.mps.is_available():
device = torch.device('mps')
shared.model = shared.model.to(device)
else:

4
modules/models.py
View File

@ -147,7 +147,7 @@ def huggingface_loader(model_name):
# Load the model in simple 16-bit mode by default
if not any([shared.args.cpu, shared.args.load_in_8bit, shared.args.load_in_4bit, shared.args.auto_devices, shared.args.disk, shared.args.deepspeed, shared.args.gpu_memory is not None, shared.args.cpu_memory is not None]):
model = LoaderClass.from_pretrained(Path(f"{shared.args.model_dir}/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.bfloat16 if shared.args.bf16 else torch.float16, trust_remote_code=shared.args.trust_remote_code)
if torch.has_mps:
if torch.backends.mps.is_available():
device = torch.device('mps')
model = model.to(device)
else:
@ -167,7 +167,7 @@ def huggingface_loader(model_name):
"trust_remote_code": shared.args.trust_remote_code
}
if not any((shared.args.cpu, torch.cuda.is_available(), torch.has_mps)):
if not any((shared.args.cpu, torch.cuda.is_available(), torch.backends.mps.is_available())):
logger.warning("torch.cuda.is_available() returned False. This means that no GPU has been detected. Falling back to CPU mode.")
shared.args.cpu = True

2
modules/text_generation.py
View File

@ -57,7 +57,7 @@ def encode(prompt, add_special_tokens=True, add_bos_token=True, truncation_lengt
return input_ids.numpy()
elif shared.args.deepspeed:
return input_ids.to(device=local_rank)
elif torch.has_mps:
elif torch.backends.mps.is_available():
device = torch.device('mps')
return input_ids.to(device)
else: