From e2c6f434ec3045cdb62a28507491cbcc9d5c53ad Mon Sep 17 00:00:00 2001
From: ShenLiang <1422485404@qq.com>
Date: Thu, 28 Nov 2019 13:11:18 +0800
Subject: [PATCH] Add Lod information for gather_nd & scatter_nd (#21404)

* add lod information, test=develop

* add lod, test=develop

* fix lod, test=develop

* fix lod, test=develop
---
 paddle/fluid/operators/gather_nd_op.cc      | 1 +
 paddle/fluid/operators/scatter_nd_add_op.cc | 1 +
 2 files changed, 2 insertions(+)

diff --git a/paddle/fluid/operators/gather_nd_op.cc b/paddle/fluid/operators/gather_nd_op.cc
index f871615d0b..3ba5beec3b 100644
--- a/paddle/fluid/operators/gather_nd_op.cc
+++ b/paddle/fluid/operators/gather_nd_op.cc
@@ -55,6 +55,7 @@ class GatherNdOp : public framework::OperatorWithKernel {
     }
 
     ctx->SetOutputDim("Out", framework::make_ddim(result_dims));
+    ctx->ShareLoD("X", /*->*/ "Out");
   }
 
  protected:
diff --git a/paddle/fluid/operators/scatter_nd_add_op.cc b/paddle/fluid/operators/scatter_nd_add_op.cc
index e86724dd56..11cf01c224 100644
--- a/paddle/fluid/operators/scatter_nd_add_op.cc
+++ b/paddle/fluid/operators/scatter_nd_add_op.cc
@@ -64,6 +64,7 @@ class ScatterNdAddOp : public framework::OperatorWithKernel {
                         "Updates has wrong shape");
     }
     ctx->SetOutputDim("Out", ref_dims);
+    ctx->ShareLoD("X", /*->*/ "Out");
   }
 
  protected:
-- 
GitLab