InferType fix (apache#2)

anijain2305 · Sep 30, 2020 · d5f9bb5 · d5f9bb5
1 parent b2c5a04
commit d5f9bb5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/python/tvm/relay/tensorrt.py b/python/tvm/relay/tensorrt.py
@@ -817,11 +817,11 @@ def EnableTrt(mod, params=None, trt_version=None, use_implicit_batch=True,
                                     transform.FoldConstant(),
                                     LegalizeLayoutTranformPass(),
                                     transform.InferType(),
-                                    # tvm.transform.PrintIR("A1"),
+                                    tvm.transform.PrintIR("A1"),
                                     transform.AnnotateTarget('tensorrt'),
-                                    # tvm.transform.PrintIR("A2"),
+                                    tvm.transform.PrintIR("A2"),
                                     transform.MergeCompilerRegions(),
-                                    # tvm.transform.PrintIR("A3"),
+                                    tvm.transform.PrintIR("A3"),
                                     transform.PartitionGraph(),
                                     transform.InferType()])
     with tvm.transform.PassContext(opt_level=3):

diff --git a/src/relay/transforms/annotate_target.cc b/src/relay/transforms/annotate_target.cc
@@ -235,8 +235,17 @@ class AnnotateTargetRewriter : public ExprRewriter {
   Expr Rewrite_(const LetNode* op, const Expr& post) final {
     auto let = Downcast<Let>(post);
 
-    auto target_n_args = AnnotateArgs({let->value, let->body});
-    auto new_expr = Let(let->var, std::get<1>(target_n_args)[0], std::get<1>(target_n_args)[1]);
+    Expr new_expr;
+    std::pair<std::string, Array<Expr>> target_n_args;
+    bool is_functional_literal = let->value.as<FunctionNode>() != nullptr;
+    // FIXME - zhi/cody - please check if this makes sense.
+    if (is_functional_literal) {
+      target_n_args = AnnotateArgs({let->body});
+      new_expr = Let(let->var, let->value, std::get<1>(target_n_args)[0]);
+    } else {
+      target_n_args = AnnotateArgs({let->value, let->body});
+      new_expr = Let(let->var, std::get<1>(target_n_args)[0], std::get<1>(target_n_args)[1]);
+    }
     op_expr_to_target_[new_expr] = std::get<0>(target_n_args);
     return std::move(new_expr);
   }