From 488e889f3d8994e44370460c44cff4b4b5a232cb Mon Sep 17 00:00:00 2001 From: phlrain Date: Thu, 11 Apr 2019 11:42:03 +0000 Subject: [PATCH] fix split infer shape; test=develop --- paddle/fluid/operators/split_op.cc | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/paddle/fluid/operators/split_op.cc b/paddle/fluid/operators/split_op.cc index d661b276bc3..9e65ef588a3 100644 --- a/paddle/fluid/operators/split_op.cc +++ b/paddle/fluid/operators/split_op.cc @@ -39,14 +39,22 @@ class SplitOp : public framework::OperatorWithKernel { if (num > 0) { int64_t in_axis_dim = in_dims[axis]; - PADDLE_ENFORCE_EQ(in_axis_dim % num, 0, - "tensor split does not result" - " in an equal division"); - size_t out_axis_dim = in_axis_dim / num; - for (size_t i = 0; i < outs_number; ++i) { - auto dim = in_dims; - dim[axis] = out_axis_dim; - outs_dims.push_back(dim); + if (ctx->IsRuntime() || in_axis_dim > 0) { + PADDLE_ENFORCE_EQ(in_axis_dim % num, 0, + "tensor split does not result" + " in an equal division"); + size_t out_axis_dim = in_axis_dim / num; + for (size_t i = 0; i < outs_number; ++i) { + auto dim = in_dims; + dim[axis] = out_axis_dim; + outs_dims.push_back(dim); + } + } else { + for (size_t i = 0; i < outs_number; ++i) { + auto dim = in_dims; + dim[axis] = -1; + outs_dims.push_back(dim); + } } } else if (sections.size() > 0) { PADDLE_ENFORCE_EQ(sections.size(), outs_number, -- GitLab