diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc index 69468f33ed07..94fc6dde6f0b 100644 --- a/src/c_api/c_api.cc +++ b/src/c_api/c_api.cc @@ -802,7 +802,7 @@ int MXDataIterGetLabel(DataIterHandle handle, NDArrayHandle *out) { // temp hack to make label 1D // TODO(tianjun) make label 1D when label_width=0 mxnet::TShape shape = db.data[1].shape(); - if (shape[1] == 1) { + if (shape.ndim() > 1 && shape[1] == 1) { *pndarray = db.data[1].Reshape(mshadow::Shape1(shape[0])); } else { *pndarray = db.data[1];