From 88f1a98e7a2b576783c4315fe43dabd93bef3387 Mon Sep 17 00:00:00 2001
From: wanghuancoder <wanghuan29@baidu.com>
Date: Fri, 8 Sep 2023 20:02:32 +0800
Subject: [PATCH] [IR] fix getfaketensorlist bug (#57063)

* fix getfaketensorlist bug

* refine
---
 paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc | 8 ++++----
 python/paddle/jit/dy2static/partial_program.py     | 4 +++-
 2 files changed, 7 insertions(+), 5 deletions(-)

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 042929dd20d..3555ebe354a 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 7679a6d4ed5..1f68cc71be5 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):
-- 
GitLab