From ffc6bad116ce78d1c68703b51d0dcfedd6bd5a84 Mon Sep 17 00:00:00 2001
From: Kittipat Virochsiri <kittipat@fb.com>
Date: Mon, 1 May 2017 12:16:08 -0700
Subject: [PATCH] Concat axis=0

Summary: Previously, the code below would go out of bound.

Reviewed By: xianjiec

Differential Revision: D4968037

fbshipit-source-id: 3760e2cddc919c45d85ac644ac3fabf72dbaf666
---
 caffe2/python/layers/concat.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/caffe2/python/layers/concat.py b/caffe2/python/layers/concat.py
index 69ed6467294..dc2fcf496e9 100644
--- a/caffe2/python/layers/concat.py
+++ b/caffe2/python/layers/concat.py
@@ -33,6 +33,13 @@ class Concat(ModelLayer):
                 "Concat expects that limited dimensions of the input tensor"
             shapes.append(list(field_type.field_type().shape))
 
+        if axis == 0:
+            self.output_schema = schema.from_blob_list(
+                input_record[0],
+                [model.net.NextScopedBlob(name + '_output')]
+            )
+            return
+
         concat_dim = 0
         for shape in shapes:
             concat_dim += shape[axis - 1]