From 38660841804714a44f590e8b1186132d4540ea4d Mon Sep 17 00:00:00 2001
From: Hongyu Liu <43953930+phlrain@users.noreply.github.com>
Date: Mon, 15 Apr 2019 16:31:54 +0800
Subject: [PATCH] Merge pull request #16797 from phlrain/fix_split

Fix split
---
 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 a05582ae0..a43bad878 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