[BACKPORT-0.6] Make sure to visit the arguments of inlined functions (#…

…5864) Co-authored-by: abergeron <abergeron@gmail.com> Co-authored-by: abergeron <abergeron@gmail.com>
apache · Jun 21, 2020 · d12e7d1 · d12e7d1
1 parent 3a671df
commit d12e7d1
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/src/relay/backend/vm/inline_primitives.cc b/src/relay/backend/vm/inline_primitives.cc
@@ -87,12 +87,22 @@ struct PrimitiveInliner : ExprMutator {
 
     if (auto func = op.as<FunctionNode>()) {
       if (func->IsPrimitive()) {
-        return CallNode::make(GetRef<Function>(func), call->args, call->attrs, call->type_args);
+        tvm::Array<Expr> call_args;
+        for (auto arg : call->args) {
+          auto new_arg = VisitExpr(arg);
+          call_args.push_back(new_arg);
+        }
+        return CallNode::make(GetRef<Function>(func), call_args, call->attrs, call->type_args);
       }
     }
 
     if (auto global = op.as<GlobalVarNode>()) {
-      return CallNode::make(GetRef<GlobalVar>(global), call->args, call->attrs, call->type_args);
+      tvm::Array<Expr> call_args;
+      for (auto arg : call->args) {
+        auto new_arg = VisitExpr(arg);
+        call_args.push_back(new_arg);
+      }
+      return CallNode::make(GetRef<GlobalVar>(global), call_args, call->attrs, call->type_args);
     }
 
     return ExprMutator::VisitExpr_(call);