From 260f66c3165ce0c48dd1514a916da6971d981578 Mon Sep 17 00:00:00 2001 From: Chandler Zuo Date: Thu, 21 Feb 2019 19:31:21 -0800 Subject: [PATCH] Fix concat dimension check bug (#17343) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17343 See [post](https://fb.workplace.com/groups/1405155842844877/permalink/2630764056950710/) Reviewed By: dzhulgakov Differential Revision: D14163001 fbshipit-source-id: 038f15d6a58b3bc31910e7bfa47c335e25739f12 --- caffe2/operators/concat_split_op.cc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/caffe2/operators/concat_split_op.cc b/caffe2/operators/concat_split_op.cc index d630f44..57ef45c 100644 --- a/caffe2/operators/concat_split_op.cc +++ b/caffe2/operators/concat_split_op.cc @@ -179,7 +179,10 @@ OPERATOR_SCHEMA(Concat) : GetDimFromOrderString( helper.GetSingleArgument("order", "NCHW")); bool add_axis = helper.GetSingleArgument("add_axis", 0) != 0; - const int canonical_axis = canonical_axis_index_(axis, in[0].dims_size()); + int adj_size = in[0].dims_size() + (add_axis ? 1 : 0); + const int canonical_axis = canonical_axis_index_(axis, adj_size); + CAFFE_ENFORCE_LT( + canonical_axis, adj_size, "Axis not in input ndim range."); CAFFE_ENFORCE_GT(in.size(), 0); vector split_shape(1, in.size()); vector out_shape(in[0].dims().begin(), in[0].dims().end()); -- 2.7.4