From 849eccb75baa93685c80f068a51e2b4f09c41c72 Mon Sep 17 00:00:00 2001
From: Weilong Wu <veyron_wu@163.com>
Date: Sun, 7 Aug 2022 20:20:39 +0800
Subject: [PATCH] [Eager] just need to run needed nodes (#44910)

* [Eager] just need to run needed nodes

* [Eager] just need to run needed nodes
---
 paddle/fluid/eager/backward.cc | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/paddle/fluid/eager/backward.cc b/paddle/fluid/eager/backward.cc
index a3c3c2718d9..b70ec78c759 100644
--- a/paddle/fluid/eager/backward.cc
+++ b/paddle/fluid/eager/backward.cc
@@ -360,11 +360,22 @@ std::vector<paddle::experimental::Tensor> RunBackward(
                 "Node's in-degree cannot be negative.",
                 next_node->name()));
 
-        if (node_in_degree_map[next_node] == 0) {
-          if (dynamic_cast<egr::GradNodeAccumulation*>(next_node)) {
-            queue.push_front(std::move(next_node));
-          } else {
-            queue.push_back(std::move(next_node));
+        if (is_general_grad) {
+          if (node_in_degree_map[next_node] == 0 &&
+              GeneralGrad::Instance().IsNeededNodes(next_node)) {
+            if (dynamic_cast<egr::GradNodeAccumulation*>(next_node)) {
+              queue.push_front(std::move(next_node));
+            } else {
+              queue.push_back(std::move(next_node));
+            }
+          }
+        } else {
+          if (node_in_degree_map[next_node] == 0) {
+            if (dynamic_cast<egr::GradNodeAccumulation*>(next_node)) {
+              queue.push_front(std::move(next_node));
+            } else {
+              queue.push_back(std::move(next_node));
+            }
           }
         }
       }
-- 
GitLab