提交 e2c6f434 编写于 作者: S ShenLiang 提交者: Yi Liu

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
上级 c0656dcb
...@@ -55,6 +55,7 @@ class GatherNdOp : public framework::OperatorWithKernel { ...@@ -55,6 +55,7 @@ class GatherNdOp : public framework::OperatorWithKernel {
} }
ctx->SetOutputDim("Out", framework::make_ddim(result_dims)); ctx->SetOutputDim("Out", framework::make_ddim(result_dims));
ctx->ShareLoD("X", /*->*/ "Out");
} }
protected: protected:
......
...@@ -64,6 +64,7 @@ class ScatterNdAddOp : public framework::OperatorWithKernel { ...@@ -64,6 +64,7 @@ class ScatterNdAddOp : public framework::OperatorWithKernel {
"Updates has wrong shape"); "Updates has wrong shape");
} }
ctx->SetOutputDim("Out", ref_dims); ctx->SetOutputDim("Out", ref_dims);
ctx->ShareLoD("X", /*->*/ "Out");
} }
protected: protected:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册