diff options
| -rw-r--r-- | tensorflow/core/grappler/optimizers/graph_rewriter.cc | 5 | ||||
| -rw-r--r-- | tensorflow/core/grappler/optimizers/model_pruner.cc | 3 |
2 files changed, 7 insertions, 1 deletions
diff --git a/tensorflow/core/grappler/optimizers/graph_rewriter.cc b/tensorflow/core/grappler/optimizers/graph_rewriter.cc index 3a3cfdf8cb..c8b100abae 100644 --- a/tensorflow/core/grappler/optimizers/graph_rewriter.cc +++ b/tensorflow/core/grappler/optimizers/graph_rewriter.cc @@ -49,7 +49,10 @@ void GraphRewriter::ForwardInputs( for (const auto& input : original_node.input()) { string input_node_name = NodeName(input); auto itr = nodes_.find(input_node_name); - CHECK(itr != nodes_.end()); + if (itr == nodes_.end()) { + // Invalid input, preserve it as is. + *new_node->add_input() = input; + } const NodeDef* input_node = itr->second; if ((input_node->device().empty() || original_node.device().empty() || input_node->device() == original_node.device()) && diff --git a/tensorflow/core/grappler/optimizers/model_pruner.cc b/tensorflow/core/grappler/optimizers/model_pruner.cc index ebc18a43b2..caf521d979 100644 --- a/tensorflow/core/grappler/optimizers/model_pruner.cc +++ b/tensorflow/core/grappler/optimizers/model_pruner.cc @@ -33,6 +33,9 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item, for (const auto& node : item.fetch) { nodes_to_preserve.insert(NodeName(node)); } + for (const auto& feed : item.feed) { + nodes_to_preserve.insert(NodeName(feed.first)); + } for (const auto& node : item.init_ops) { nodes_to_preserve.insert(NodeName(node)); } |