fix

Author: gigiblender

src/relax/usmp/transform/assign_pool_info.cc

@@ -244,7 +244,7 @@ IRModule PoolInfoAssigner::operator()() {
           RelaxPoolInfoAssigner(func, target_pool_infos_, target_const_pool_infos_);
       relax::Expr body = relax_pool_info_assigner();
       Function new_relax_func =
-          Function(func->params, body, func->ret_type, func->attrs, func->span);
+          Function(func->params, body, func->ret_type, func->ret_shape, func->attrs, func->span);
       mod_->Update(gv, new_relax_func);
     } else if (kv.second->IsInstance<tir::PrimFuncNode>()) {
       using TIRPoolInfoAssigner = tir::usmp::TIRPoolInfoAssigner;