提交 5c131d2a 编写于 作者: M Megvii Engine Team

fix(device): fix get_device_count for xpu

GitOrigin-RevId: 6c77f507f3ef4436fb9674f988b2e734af475612
上级 9eb4ba88
......@@ -44,7 +44,9 @@ def _str2device_type(type_str: str, allow_unspec: bool = True):
elif type_str == "ROCM" or type_str == "AMDGPU":
return DeviceType.ROCM
else:
assert allow_unspec and str == "XPU", "device type can only be cpu, gpu or xpu"
assert (
allow_unspec and type_str == "XPU"
), "device type can only be cpu, gpu or xpu"
return DeviceType.UNSPEC
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册