Experimental Intel Arc transformers support (untested)

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

View File

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

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():
device = torch.device('mps')
return input_ids.to(device)
elif hasattr(torch, 'xpu') and torch.xpu.is_available():
return input_ids.to('xpu')
else:
return input_ids.cuda()