1 files changed, 10 insertions, 5 deletions
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_reorder_axes.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_reorder_axes.cc
index 0d274fc687..4a264e1cf1 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/convert_reorder_axes.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_reorder_axes.cc
@@ -86,9 +86,12 @@ TransposeOperator* CreateTransposeFromReorderAxes(
 
 // Converts ReorderAxes into Transpose and Reshape which are compatible with the
 // TFLite interpreter.
-bool ConvertReorderAxes::Run(Model* model, std::size_t op_index) {
+::tensorflow::Status ConvertReorderAxes::Run(Model* model, std::size_t op_index,
+                                             bool* modified) {
+  *modified = false;
   auto reorder_it = model->operators.begin() + op_index;
-  if (reorder_it->get()->type != OperatorType::kReorderAxes) return false;
+  if (reorder_it->get()->type != OperatorType::kReorderAxes)
+    return ::tensorflow::Status::OK();
 
   auto* reorder_op = static_cast<ReorderAxesOperator*>(reorder_it->get());
   CHECK_EQ(reorder_op->inputs.size(), 1);
@@ -113,8 +116,9 @@ bool ConvertReorderAxes::Run(Model* model, std::size_t op_index) {
   // Yield if input array contains constants or if output array size has not
   // been adjusted to reflect the permutations in ReorderAxes. ReorderAxes will
   // be merged into a constant array when possible.
-  if (IsConstantParameterArray(*model, constant_input_array_name)) return false;
-  if (!output_array.has_shape()) return false;
+  if (IsConstantParameterArray(*model, constant_input_array_name))
+    return ::tensorflow::Status::OK();
+  if (!output_array.has_shape()) return ::tensorflow::Status::OK();
 
   const auto input_axes_order = reorder_op->input_axes_order;
   const auto output_axes_order = reorder_op->output_axes_order;
@@ -143,7 +147,8 @@ bool ConvertReorderAxes::Run(Model* model, std::size_t op_index) {
   CHECK_EQ(reorder_it->get(), reorder_op);
   model->operators.erase(reorder_it);
 
-  return true;
+  *modified = true;
+  return ::tensorflow::Status::OK();
 }
 
 }  // namespace toco