Fix infer shape of split lod tensor op

test=develop

Fix infer shape of split lod tensor op
test=develop
592011bb · minqiyang · 08e75731 · 592011bb
显示空白变更内容
内联并排

Showing with 3 addition and 1 deletion

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

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

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

    context->SetOutputDim("OutTrue", context->GetInputDim("X"));
    context->SetOutputDim("OutFalse", context->GetInputDim("X"));