Merge pull request #16852 from sneaxiy/fix_merge_lod_tensor_op_infer_shape

Fix merge_lod_tensor_op infer shape

Merge pull request #16852 from sneaxiy/fix_merge_lod_tensor_op_infer_shape
Fix merge_lod_tensor_op infer shape
0cc984bd · Tao Luo · phlrain · 87be315c · 0cc984bd
隐藏空白更改
内联并排

Showing with 3 addition and 1 deletion

paddle/fluid/operators/merge_lod_tensor_op.cc paddle/fluid/operators/merge_lod_tensor_op.cc +3 -1

未找到文件。
--- a/paddle/fluid/operators/merge_lod_tensor_op.cc
+++ b/paddle/fluid/operators/merge_lod_tensor_op.cc
@@ -164,7 +164,9 @@ class MergeLoDTensorInferShape : public framework::InferShapeBase {

    auto mask_dim = context->GetInputDim("Mask");
    PADDLE_ENFORCE_EQ(mask_dim.size(), 2);
-    PADDLE_ENFORCE_EQ(mask_dim[1], 1);
+    if (context->IsRuntime() || mask_dim[1] > 0) {
+      PADDLE_ENFORCE_EQ(mask_dim[1], 1);
+    }

    context->SetOutputDim("Out", context->GetInputDim("InTrue"));
  }