Convert Squeeze into Reshape: Support empty output shapes.

PiperOrigin-RevId: 188740288

diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_squeeze_to_reshape.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_squeeze_to_reshape.cc
index e601284..81cedb5 100644 (file)
--- a/tensorflow/contrib/lite/toco/graph_transformations/convert_squeeze_to_reshape.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_squeeze_to_reshape.cc
@@ -57,6 +57,11 @@ bool ConvertSqueezeToReshape::Run(Model* model, std::size_t op_index) {
   // We use the output shape that has been calculated by shape propagation.
   const auto& output_shape = model->GetArray(squeeze_op->outputs[0]).shape();
 
+  // Empty shapes will not work as empty data arrays.
+  if (output_shape.dimensions_count() == 0) {
+    return false;
+  }
+
   auto* reshape_op = new TensorFlowReshapeOperator;
   reshape_op->inputs = {
       squeeze_op->inputs[0],