From 5c131d2a74527393234809dd24a07ae7972fa735 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Tue, 30 Mar 2021 16:46:28 +0800
Subject: [PATCH] fix(device): fix get_device_count for xpu

GitOrigin-RevId: 6c77f507f3ef4436fb9674f988b2e734af475612
---
 imperative/python/megengine/device.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/imperative/python/megengine/device.py b/imperative/python/megengine/device.py
index 63faa41bb..8a31754ae 100644
--- a/imperative/python/megengine/device.py
+++ b/imperative/python/megengine/device.py
@@ -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
 
 
-- 
GitLab