[Releay] Fix on_device call for explicit virtual_device (#12088)

apache · Jul 14, 2022 · 2fee86e · 2fee86e
1 parent 484f602
commit 2fee86e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/tvm/relay/op/annotation/annotation.py b/python/tvm/relay/op/annotation/annotation.py
@@ -28,6 +28,8 @@ def _make_virtual_device(device):
         return target.VirtualDevice(device)
     if isinstance(device, str):
         return target.VirtualDevice(_nd.device(device))
+    if isinstance(device, target.VirtualDevice):
+        return device
     raise ValueError("expecting a Device or device name, but received a %s" % (type(device)))