diff --git a/__init__.py b/__init__.py index bbccd67..d1e710a 100644 --- a/__init__.py +++ b/__init__.py @@ -47,6 +47,11 @@ def get_device_list(): devs += [f"xpu:{i}" for i in range(torch.xpu.device_count())] except Exception: pass + try: + if torch.backends.mps.is_available(): + devs += ["mps"] + except Exception: + pass return devs def set_current_device(device):