Fix optimize

apache · Feb 12, 2020 · aaf62e4 · aaf62e4
1 parent 176ffe5
commit aaf62e4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/relay/build_module.py b/python/tvm/relay/build_module.py
@@ -275,7 +275,7 @@ def optimize(mod, target=None, params=None):
     params : dict
         The parameters of the final graph.
     """
-    if isinstance(mod, _Module):
+    if isinstance(mod, IRModule):
         func = mod["main"]
     elif isinstance(mod, _expr.Function):
         func = mod