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
6 years ago · e2c6f434ec
parent c0656dcb1a
commit e2c6f434ec
2 changed files with 2 additions and 0 deletions
--- 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:
--- 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: