diff --git a/src/relay/backend/vm/inline_primitives.cc b/src/relay/backend/vm/inline_primitives.cc index 6a7962697101..1d6ba4a46c61 100644 --- 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()) { if (func->IsPrimitive()) { - return CallNode::make(GetRef(func), call->args, call->attrs, call->type_args); + tvm::Array call_args; + for (auto arg : call->args) { + auto new_arg = VisitExpr(arg); + call_args.push_back(new_arg); + } + return CallNode::make(GetRef(func), call_args, call->attrs, call->type_args); } } if (auto global = op.as()) { - return CallNode::make(GetRef(global), call->args, call->attrs, call->type_args); + tvm::Array call_args; + for (auto arg : call->args) { + auto new_arg = VisitExpr(arg); + call_args.push_back(new_arg); + } + return CallNode::make(GetRef(global), call_args, call->attrs, call->type_args); } return ExprMutator::VisitExpr_(call);