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

Experimental Intel Arc transformers support (untested)

Browse source
This commit is contained in:
oobabooga 2023-10-15 20:51:11 -07:00
parent 8cce1f1126
commit b88b2b74a6
2 changed files with 4 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
modules/models.py
View file

@ -137,6 +137,8 @@ def huggingface_loader(model_name):
if torch.backends.mps.is_available(): if torch.backends.mps.is_available():
device = torch.device('mps') device = torch.device('mps')
model = model.to(device) model = model.to(device)
elif hasattr(torch, 'xpu') and torch.xpu.is_available():
model = model.to('xpu')
else: else:
model = model.cuda() model = model.cuda()

2
modules/text_generation.py
View file

@ -132,6 +132,8 @@ def encode(prompt, add_special_tokens=True, add_bos_token=True, truncation_lengt
elif torch.backends.mps.is_available(): elif torch.backends.mps.is_available():
device = torch.device('mps') device = torch.device('mps')
return input_ids.to(device) return input_ids.to(device)
elif hasattr(torch, 'xpu') and torch.xpu.is_available():
return input_ids.to('xpu')
else: else:
return input_ids.cuda() return input_ids.cuda()