提交 2797fcfa 编写于 作者: M Megvii Engine Team

fix(mge/device): add missed API to __all__ scope

GitOrigin-RevId: c3ce90990f09ca4a6d00a58e7e1b89ecaa059288
上级 d7c546c9
......@@ -14,6 +14,9 @@ from .core._imperative_rt.utils import _try_coalesce_all_free_memory
__all__ = [
"is_cuda_available",
"is_cambricon_available",
"is_atlas_available",
"is_rocm_available",
"get_device_count",
"get_default_device",
"set_default_device",
......@@ -30,7 +33,6 @@ __all__ = [
"reset_max_memory_stats",
"set_prealloc_config",
"coalesce_free_memory",
"DeviceType",
]
......@@ -252,6 +254,7 @@ def set_prealloc_config(
def what_is_xpu():
r"""Return the precise device type like ``cpu``, ``cuda`` and so on."""
return _what_is_xpu().name.lower()
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册