Do not execute _SOURCE and _SINK nodes on STE.
PiperOrigin-RevId: 442567172
diff --git a/tensorflow/core/common_runtime/single_threaded_executor.cc b/tensorflow/core/common_runtime/single_threaded_executor.cc
index c2f7a26..faf810b 100644
--- a/tensorflow/core/common_runtime/single_threaded_executor.cc
+++ b/tensorflow/core/common_runtime/single_threaded_executor.cc
@@ -94,16 +94,21 @@
ordered_nodes.size());
}
- kernels_.reserve(ordered_nodes.size());
+ // We reserve two less nodes because we do not need to create kernels for
+ // the _SOURCE and _SINK nodes.
+ kernels_.reserve(ordered_nodes.size() - 2);
std::vector<Node*> nodes_with_kernels;
std::vector<Node*> nodes_with_const_tensor_kernels;
- nodes_with_kernels.reserve(ordered_nodes.size());
+ nodes_with_kernels.reserve(ordered_nodes.size() - 2);
std::map<size_t, Node*> arg_index_to_node_map;
absl::flat_hash_map<Node*, size_t> node_to_index_map;
// Create the kernel and input-related structures for each node in `graph`.
for (Node* n : ordered_nodes) {
+ if (n->IsSource() || n->IsSink()) {
+ continue;
+ }
TF_RETURN_IF_ERROR(ValidateOpIsSafeForSyncExecution(
*n, params_.allow_control_flow_sync_execution));
if (n->IsArg()) {
