fix_test_write_python_container (#61692)

huangjiyi · web-flow · commit a0f9438a58f6 · 2024-02-08T14:49:08.000+08:00
diff --git a/paddle/fluid/framework/new_executor/instruction/control_flow/while_instruction.cc b/paddle/fluid/framework/new_executor/instruction/control_flow/while_instruction.cc
@@ -209,20 +209,11 @@ void WhileInstruction::ShareDatasToOutputs() {
     auto& out_var_name = body_outputs_[i + 1];
     auto* out_var = body_inter_->local_scope()->GetVar(out_var_name);
     VLOG(6) << "share data from " << out_var_name << " -> " << i << " output";
-
     if (out_var->IsType<phi::DenseTensor>()) {
       outputs_[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
           out_var->Get<phi::DenseTensor>());
       VLOG(6) << "share data from " << out_var_name << "[" << out_var << "]"
               << " -> " << i << " output[" << outputs_[i] << "]";
-
-      // NOTE(zhangbo): Delete the input of the yield operator, except for the
-      // external vars of the block.
-      if (external_input_names_.count(out_var_name) == 0) {
-        VLOG(6) << "clear internel input " << out_var_name;
-        out_var->GetMutable<phi::DenseTensor>()->clear();
-      }
-
     } else if (out_var->IsType<phi::TensorArray>()) {
       const auto& inner_array = out_var->Get<phi::TensorArray>();
       auto* output_array = outputs_[i]->GetMutable<phi::TensorArray>();
@@ -234,6 +225,19 @@ void WhileInstruction::ShareDatasToOutputs() {
 
     VLOG(6) << "done";
   }
+
+  for (size_t i = 0; i < outputs_.size(); ++i) {
+    auto& out_var_name = body_outputs_[i + 1];
+    auto* out_var = body_inter_->local_scope()->GetVar(out_var_name);
+    if (out_var->IsType<phi::DenseTensor>()) {
+      // NOTE(zhangbo): Delete the input of the yield operator, except for the
+      // external vars of the block.
+      if (external_input_names_.count(out_var_name) == 0) {
+        VLOG(6) << "clear internel input " << out_var_name;
+        out_var->GetMutable<phi::DenseTensor>()->clear();
+      }
+    }
+  }
 }
 
 void WhileInstruction::Run() {
diff --git a/test/dygraph_to_static/test_write_python_container.py b/test/dygraph_to_static/test_write_python_container.py
@@ -127,6 +127,7 @@ def test_write_container_sot(self):
         self.assertEqual(out_static, out_dygraph)
 
     @test_ast_only
+    @test_legacy_and_pt_and_pir
     def test_write_container(self):
         func_static = paddle.jit.to_static(self.func)
         input = paddle.to_tensor([1, 2, 3])
