diff --git a/tensorflow_serving/servables/tensorflow/predict_util.cc b/tensorflow_serving/servables/tensorflow/predict_util.cc
index ca3d62611c6..6728682d78f 100644
--- a/tensorflow_serving/servables/tensorflow/predict_util.cc
+++ b/tensorflow_serving/servables/tensorflow/predict_util.cc
@@ -111,6 +111,13 @@ Status PreProcessPrediction(const SignatureDef& signature,
}
inputs->emplace_back(std::make_pair(iter->second.name(), tensor));
}
+ // Sort inputs to make TF session not insert too many keys to executor
+ std::sort(
+ inputs->begin(), inputs->end(),
+ [](std::pair &left, std::pair &right) {
+ return left.first < right.first;
+ }
+ );
// Prepare run target.
std::set seen_outputs;
@@ -134,6 +141,7 @@ Status PreProcessPrediction(const SignatureDef& signature,
output_tensor_names->emplace_back(iter->second.name());
output_tensor_aliases->emplace_back(alias);
}
+
// When no output is specified, fetch all output tensors specified in
// the signature.
if (output_tensor_names->empty()) {
@@ -142,6 +150,7 @@ Status PreProcessPrediction(const SignatureDef& signature,
output_tensor_aliases->emplace_back(iter.first);
}
}
+
return Status::OK();
}