Use main variable instead of entry_func

apache · Jul 7, 2019 · 3be301d · 3be301d
1 parent ee4d514
commit 3be301d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py
@@ -94,7 +94,7 @@ def optimize(self, prog: Expr):
         opts = relay.transform.Sequential([relay.transform.SimplifyInference(),
                                            relay.transform.FuseOps(fuse_opt_level=0)])
         mod = opts(mod)
-        optimized = mod[mod.entry_func]
+        optimized = mod['main']
         return optimized if isinstance(unwrapped, Function) else optimized.body