Fix test case

apache · Aug 22, 2023 · 685cc76 · 685cc76
1 parent 28d0cc2
commit 685cc76
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tests/python/relax/test_relay_translator.py b/tests/python/relax/test_relay_translator.py
@@ -330,10 +330,12 @@ def run_after_pass(self, mod, info):
             self.events.append("run after " + info.name)
 
     my_test = SampleRunBeforeAfterInstrument()
-    with tvm.transform.PassContext(instruments=[my_test]):
-        mod = relay.transform.InferType()(mod)
+    relax_mod_with_attrs = relay_translator.from_relay(
+        mod["main"], target="llvm", instruments=[my_test]
+    )
 
-    assert "run before InferType" "run after InferType" == "".join(my_test.events)
+    assert "run after " in "".join(my_test.events)
+    assert "run before " in "".join(my_test.events)
 
 
 if __name__ == "__main__":