diff options
Diffstat (limited to 'tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc')
| -rw-r--r-- | tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc index adfae2e1a3..f46c30c92c 100644 --- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc @@ -2233,6 +2233,9 @@ string ArithmeticOptimizer::TrySimplifyAndReplaceUses( new_square_node->set_input(i - 1, new_square_node->input(i)); } new_square_node->mutable_input()->RemoveLast(); + for (const string& input : new_square_node->input()) { + node_map_->AddOutput(NodeName(input), new_square_node->name()); + } return new_square_node->name(); } } |