Improve style

apache · Nov 9, 2019 · a581dc8 · a581dc8
1 parent 4c6ec7d
commit a581dc8
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
@@ -453,14 +453,11 @@ def get_name(node):
 
 def infer_type(node, mod=None):
     """A method to infer the type of an intermediate node in the relay graph."""
-    if isinstance(mod, _module.Module):
-        mod["main"] = _expr.Function([], node)
-        mod = _transform.InferType()(mod)
-        entry = mod["main"]
-        return entry.body
-    mod = _module.Module.from_expr(node)
-    mod = _transform.InferType()(mod)
-    entry = mod["main"]
+    new_mod = _module.Module.from_expr(node)
+    if mod is not None:
+        new_mod.update(mod)
+    new_mod = _transform.InferType()(new_mod)
+    entry = new_mod["main"]
     return entry if isinstance(node, _expr.Function) else entry.body
 
 def infer_shape(inputs, mod=None):