Make sure to visit the arguments of inlined functions (#4783)

apache · Jan 30, 2020 · 6798ba8 · 6798ba8
1 parent 1b8522e
commit 6798ba8
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);