"...git@developer.sourcefind.cn:chenpangpang/open-webui.git" did not exist on "8dacc86ab0156307ed304996f5fadae24bf21864"
Unverified Commit 8c00b53e authored by Felix Hirwa Nshuti's avatar Felix Hirwa Nshuti Committed by GitHub
Browse files

Adding `mps` as device for `Pipeline` class (#30080)



* adding env variable for mps and is_torch_mps_available for Pipeline

* fix linting errors

* Remove environment overide
Co-authored-by: default avataramyeroberts <22614925+amyeroberts@users.noreply.github.com>

---------
Co-authored-by: default avataramyeroberts <22614925+amyeroberts@users.noreply.github.com>
parent 7afade20
...@@ -42,6 +42,7 @@ from ..utils import ( ...@@ -42,6 +42,7 @@ from ..utils import (
is_torch_available, is_torch_available,
is_torch_cuda_available, is_torch_cuda_available,
is_torch_mlu_available, is_torch_mlu_available,
is_torch_mps_available,
is_torch_npu_available, is_torch_npu_available,
is_torch_xpu_available, is_torch_xpu_available,
logging, logging,
...@@ -860,6 +861,8 @@ class Pipeline(_ScikitCompat): ...@@ -860,6 +861,8 @@ class Pipeline(_ScikitCompat):
self.device = torch.device(f"npu:{device}") self.device = torch.device(f"npu:{device}")
elif is_torch_xpu_available(check_device=True): elif is_torch_xpu_available(check_device=True):
self.device = torch.device(f"xpu:{device}") self.device = torch.device(f"xpu:{device}")
elif is_torch_mps_available():
self.device = torch.device(f"mps:{device}")
else: else:
raise ValueError(f"{device} unrecognized or not available.") raise ValueError(f"{device} unrecognized or not available.")
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment