diff options
-rw-r--r-- | tensorflow/contrib/lite/toco/tooling_util.cc | 15 |
1 files changed, 10 insertions, 5 deletions
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc index f3daac175e..39b030c338 100644 --- a/tensorflow/contrib/lite/toco/tooling_util.cc +++ b/tensorflow/contrib/lite/toco/tooling_util.cc @@ -1138,11 +1138,16 @@ void ResolveModelFlags(const ModelFlags& model_flags, Model* model) { } } } else { - const auto& input_array_dims = - *input_array.mutable_shape()->mutable_dims(); - CHECK_EQ(input_array_dims.size(), input_array_proto.shape().dims_size()); - for (int i = 0; i < input_array_dims.size(); i++) { - CHECK_EQ(input_array_dims[i], input_array_proto.shape().dims(i)); + if (input_array_proto.has_shape()) { + // If an input shape was specified on the flags ensure that it matches + // the actual shape in the model. + const auto& input_array_dims = + *input_array.mutable_shape()->mutable_dims(); + CHECK_EQ(input_array_dims.size(), + input_array_proto.shape().dims_size()); + for (int i = 0; i < input_array_dims.size(); i++) { + CHECK_EQ(input_array_dims[i], input_array_proto.shape().dims(i)); + } } } |