OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 37 Packages Projects Releases Wiki Activity

Concat axis=0

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

Reviewed By: xianjiec

Differential Revision: D4968037

fbshipit-source-id: 3760e2cddc919c45d85ac644ac3fabf72dbaf666
This commit is contained in:
Kittipat Virochsiri 2017-05-01 12:16:08 -07:00 committed by Facebook Github Bot
parent 1040b5f91c
commit ffc6bad116
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
caffe2/python/layers/concat.py
View File

@ -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]