diff --git a/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc
index 042929dd20d2d23ce2acbf6c802a00b3af40e562..3555ebe354ab7d741ed0131d6ceb604e26ef12b8 100644
--- a/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc
+++ b/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc
@@ -359,12 +359,12 @@ std::vector<std::shared_ptr<phi::TensorBase>> GetFakeTensorList(
   } else if (input_type.isa<ir::VectorType>()) {
     auto vec_inner_types = input_type.dyn_cast<ir::VectorType>().data();
     for (size_t i = 0; i < vec_inner_types.size(); ++i) {
-      if (vec_inner_types[0].isa<dialect::AllocatedDenseTensorType>()) {
+      if (vec_inner_types[i].isa<dialect::AllocatedDenseTensorType>()) {
         vec_res.push_back(build_fake_dense_tensor(
-            vec_inner_types[0].dyn_cast<dialect::AllocatedDenseTensorType>()));
-      } else if (vec_inner_types[0].isa<dialect::AllocatedSelectedRowsType>()) {
+            vec_inner_types[i].dyn_cast<dialect::AllocatedDenseTensorType>()));
+      } else if (vec_inner_types[i].isa<dialect::AllocatedSelectedRowsType>()) {
         vec_res.push_back(build_fake_selected_rows(
-            vec_inner_types[0].dyn_cast<dialect::AllocatedSelectedRowsType>()));
+            vec_inner_types[i].dyn_cast<dialect::AllocatedSelectedRowsType>()));
       }
     }
   }
diff --git a/python/paddle/jit/dy2static/partial_program.py b/python/paddle/jit/dy2static/partial_program.py
index 7679a6d4ed513e9f103be9ed2c0e5c8881aebdbc..1f68cc71be5ff21acc7d389eecd8714471cb6796 100644
--- a/python/paddle/jit/dy2static/partial_program.py
+++ b/python/paddle/jit/dy2static/partial_program.py
@@ -635,7 +635,9 @@ class PartialProgramLayer:
             filter(_need_aggregation, self._outputs.tolist())
         )
         for _var in to_processed_vars:
-            _insert_aggregation_ops_for_var(target_program, _var)
+            target_program: paddle.static.Program
+            target_var = target_program.global_block().var(_var.name)
+            _insert_aggregation_ops_for_var(target_program, target_var)
 
     @switch_to_static_graph
     def _append_backward_desc(self, main_program):