未验证 提交 19511dfa 编写于 作者: L liym27 提交者: GitHub

Fix bug in assign op: support to infer shape for LOD_TENSOR_ARRAY. (#24268)

上级 b43b46e6
...@@ -36,6 +36,13 @@ class AssignOp : public framework::OperatorWithKernel { ...@@ -36,6 +36,13 @@ class AssignOp : public framework::OperatorWithKernel {
if (type == framework::proto::VarType::LOD_TENSOR) { if (type == framework::proto::VarType::LOD_TENSOR) {
ctx->ShareLoD("X", /*->*/ "Out"); ctx->ShareLoD("X", /*->*/ "Out");
} }
} else if (type == framework::proto::VarType::LOD_TENSOR_ARRAY) {
if (ctx->IsRuntime()) {
// The runtime output shape is determined in kernel.
return;
} else {
ctx->SetOutputDim("Out", ctx->GetInputDim("X"));
}
} }
} }
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册