diff --git a/python/tvm/relay/build_module.py b/python/tvm/relay/build_module.py
index fa812cb35703..8b347883fe3a 100644
--- 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