Merge pull request #16852 from sneaxiy/fix_merge_lod_tensor_op_infer_shape

luotao1 · phlrain · commit 0cc984bdb258 · 2019-04-17T02:08:40.000Z
Fix merge_lod_tensor_op infer shape
diff --git 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"));
   }

Original file line number	Diff line number	Diff line change
`@@ -164,7 +164,9 @@ class MergeLoDTensorInferShape : public framework::InferShapeBase {`
`164`	`164`
`165`	`165`	`auto mask_dim = context->GetInputDim("Mask");`
`166`	`166`	`PADDLE_ENFORCE_EQ(mask_dim.size(), 2);`
`167`		`- PADDLE_ENFORCE_EQ(mask_dim[1], 1);`
	`167`	`+ if (context->IsRuntime() \|\| mask_dim[1] > 0) {`
	`168`	`+ PADDLE_ENFORCE_EQ(mask_dim[1], 1);`
	`169`	`+ }`
`168`	`170`
`169`	`171`	`context->SetOutputDim("Out", context->GetInputDim("InTrue"));`
`170`	`172`	`}`