From 3228fc34c91aef1f24edfd563d3238f158e0d607 Mon Sep 17 00:00:00 2001
From: 0x45f <23097963+0x45f@users.noreply.github.com>
Date: Fri, 25 Mar 2022 10:50:16 +0800
Subject: [PATCH] Fix loop index for FillZeroForEmptyGradInputs (#40909)

* Fix loop index for FillZeroForEmptyGradInputs

* Call fill zero in run_program_grad
---
 paddle/fluid/eager/auto_code_generator/eager_generator.cc       | 2 +-
 .../auto_code_generator/final_state_generator/codegen_utils.py  | 2 +-
 paddle/fluid/eager/to_static/run_program_op_node.h              | 1 +
 paddle/fluid/eager/utils.cc                                     | 2 +-
 4 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/eager/auto_code_generator/eager_generator.cc b/paddle/fluid/eager/auto_code_generator/eager_generator.cc
index fc3ea23492..0be4be6c96 100644
--- a/paddle/fluid/eager/auto_code_generator/eager_generator.cc
+++ b/paddle/fluid/eager/auto_code_generator/eager_generator.cc
@@ -48,7 +48,7 @@ static std::unordered_map<std::string, paddle::framework::AttributeMap>
     operators_with_attrs = {};
 
 static std::unordered_set<std::string> ops_to_fill_zero_for_empty_grads = {
-    "split"};
+    "split", "rnn"};
 
 /* --- Black Ops list that's NO NEED to apply code generation --- */
 static std::unordered_set<std::string> black_ops_list = {"run_program"};
diff --git a/paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py b/paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py
index e907b6bd1b..56cbc05b1a 100644
--- a/paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py
+++ b/paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py
@@ -21,7 +21,7 @@ import os
 ########################
 ### Global Variables ###
 ########################
-ops_to_fill_zero_for_empty_grads = set(list("split"))
+ops_to_fill_zero_for_empty_grads = set(["split", "rnn"])
 
 # For API dispatch used at python-level
 # { op_name : [arg_name, ...] }
diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h
index a60d7b5c65..3c2fa59cf1 100644
--- a/paddle/fluid/eager/to_static/run_program_op_node.h
+++ b/paddle/fluid/eager/to_static/run_program_op_node.h
@@ -367,6 +367,7 @@ class GradNodeRunProgram : public egr::GradNodeBase {
         paddle::platform::errors::InvalidArgument(
             "The out_grads.size() of RunProgramGradOp should be equal to 1."));
 
+    egr::EagerUtils::FillZeroForEmptyGradInputs(&grads, this->InputMeta());
     VLOG(3) << "out_grads[0].size() : " << grads[0].size();
     std::vector<paddle::experimental::Tensor> x_grad;
     std::vector<paddle::experimental::Tensor> params_grad;
diff --git a/paddle/fluid/eager/utils.cc b/paddle/fluid/eager/utils.cc
index f25c4dfcd5..aee7c87361 100644
--- a/paddle/fluid/eager/utils.cc
+++ b/paddle/fluid/eager/utils.cc
@@ -398,7 +398,7 @@ void EagerUtils::FillZeroForEmptyGradInputs(
     std::vector<std::vector<paddle::experimental::Tensor>>* in_grads,
     const std::vector<std::vector<GradSlotMeta>>& grad_in_metas) {
   for (size_t i = 0; i < in_grads->size(); i++) {
-    for (size_t j = 0; j < (*in_grads)[0].size(); j++) {
+    for (size_t j = 0; j < (*in_grads)[i].size(); j++) {
       paddle::experimental::Tensor& grad = (*in_grads)[i][j];
       if (!grad.is_initialized()) {
         const GradSlotMeta& grad_in_meta = grad_in_metas[i][j];
-- 
GitLab