Fix serialization tests

apache · Jan 6, 2022 · 592cf52 · 592cf52
1 parent 8cc9930
commit 592cf52
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/relay/ir/expr.cc b/src/relay/ir/expr.cc
@@ -149,7 +149,7 @@ Var WithFields(Var var, Optional<Id> opt_vid, Optional<Type> opt_type_annotation
     VarNode* cow_var_node = var.CopyOnWrite();
     cow_var_node->vid = vid;
     cow_var_node->type_annotation = type_annotation;
-    cow_var_node->virtual_device_ = var->virtual_device();
+    cow_var_node->virtual_device_ = virtual_device;
     cow_var_node->span = span;
   }
   return var;

diff --git a/tests/python/relay/test_json_compact.py b/tests/python/relay/test_json_compact.py
@@ -49,7 +49,7 @@ def test_var():
         {"type_key": ""},
         {
             "type_key": "relay.Var",
-            "attrs": {"_checked_type_": "0", "span": "0", "type_annotation": "0", "vid": "2"},
+            "attrs": {"_checked_type_": "0", "span": "0", "type_annotation": "0", "vid": "2", "virtual_device_": "0"},
         },
         {"type_key": "relay.Id", "attrs": {"name_hint": "a3"}},
         {"type_key": "relay.TensorType", "attrs": {"dtype": "float32", "shape": "4", "span": "0"}},
@@ -117,7 +117,7 @@ def test_global_var():
         {"type_key": ""},
         {
             "type_key": "relay.GlobalVar",
-            "attrs": {"_checked_type_": "0", "name_hint": "x", "span": "0"},
+            "attrs": {"_checked_type_": "0", "name_hint": "x", "span": "0", "virtual_device_": "0"},
         },
     ]
     data = {
@@ -130,7 +130,7 @@ def test_global_var():
     assert isinstance(tvar, tvm.ir.GlobalVar)
     nodes = [
         {"type_key": ""},
-        {"type_key": "GlobalVar", "attrs": {"_checked_type_": "0", "name_hint": "x", "span": "0"}},
+        {"type_key": "GlobalVar", "attrs": {"_checked_type_": "0", "name_hint": "x", "span": "0", "virtual_device_": "0"}},
     ]
     data = {
         "root": 1,