Fix optimize

aaf62e47 · tqchen · Tianqi Chen · 176ffe50 · aaf62e47
Commit aaf62e47 authored Feb 11, 2020 by tqchen Committed by Tianqi Chen Feb 12, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 additions and 1 deletions

python/tvm/relay/build_module.py
+1 -1

No files found.
--- 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