 
     memory::allocation::AllocationPtr tmp_dev_ins_data;
     const T** dev_ins_data = nullptr;
-    if (!has_same_shape || (in_num != 2)) {
+    if (!has_same_shape || in_num < 2 || in_num > 4) {
       tmp_dev_ins_data =
           platform::DeviceTemporaryAllocator::Instance().Get(context).Allocate(
               inputs_data.size() * sizeof(T*));

        SplitKernel<<<grid_dims, block_dims, 0, context.stream()>>>(
            input.data<T>(), in_row, in_col, out0_col, outputs_data[0],
            outputs_data[1]);
      } else {
        SplitKernel<<<grid_dims, block_dims, 0, context.stream()>>>(
            input.data<T>(), in_row, in_col, out0_col, dev_out_gpu_data);


    memory::allocation::AllocationPtr tmp_dev_outs_data;
    T** dev_out_gpu_data = nullptr;
    if (!has_same_shape || o_num < 2 || o_num > 4) {
      tmp_dev_outs_data =
          platform::DeviceTemporaryAllocator::Instance().Get(context).Allocate(
              outputs_data.size() * sizeof(T*));
