Fix latent bug in dependency optimizer.

PiperOrigin-RevId: 184291701

diff --git a/tensorflow/core/grappler/optimizers/dependency_optimizer.cc b/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
index 0842fc9..7b4ca14 100644 (file)
--- a/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
@@ -97,7 +97,12 @@ bool DependencyOptimizer::SafeToRemoveIdentity(const NodeDef& node) {
       // TODO(rmlarsen): Try to remove this artificial contraint.
       return false;
     }
-    for (auto consumer : node_map_->GetOutputs(node.name())) {
+  }
+  for (auto consumer : node_map_->GetOutputs(node.name())) {
+    if (node.input_size() > 1 && IsMerge(*consumer)) {
+      return false;
+    }
+    if (IsSwitch(*input)) {
       for (const string& consumer_input : consumer->input()) {
         if (consumer_input == AsControlDependency(node.name())) {
           return false;