From 2863d44374c226444785f81ffe385e5dd5b9c295 Mon Sep 17 00:00:00 2001 From: Yuefeng Zhou Date: Fri, 2 Feb 2018 12:46:34 -0800 Subject: [PATCH] Skip the node that has unexpected number of outputs. PiperOrigin-RevId: 184320865 --- tensorflow/core/graph/costmodel.cc | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/graph/costmodel.cc b/tensorflow/core/graph/costmodel.cc index f47c983..4f3a6ec 100644 --- a/tensorflow/core/graph/costmodel.cc +++ b/tensorflow/core/graph/costmodel.cc @@ -252,9 +252,12 @@ void CostModel::RecordMaxMemorySize(const Node* node, int output_slot, const DataType& dtype) { const int id = Id(node); if (id < 0) return; - CHECK_LT(output_slot, node->num_outputs()) - << "Unexpected output slot for node " << node->DebugString() << ". Got " - << output_slot << " but its num_outputs is " << node->num_outputs(); + if (output_slot >= node->num_outputs()) { + LOG(ERROR) << "Unexpected output slot for node " << node->DebugString() + << ". Got " << output_slot << " but its num_outputs is " + << node->num_outputs(); + return; + } Ensure(id, node->num_outputs()); auto& current_max = max_mem_usage_[id].output_port_mem[output_slot]; // If the memory allocator doesn't track memory usage, let's infer a lower -- 2.7.4