diff --git a/tensorflow/lite/delegates/gpu/cl/inference_context.cc b/tensorflow/lite/delegates/gpu/cl/inference_context.cc index e0ad1355cd1..10a502040bf 100644 --- a/tensorflow/lite/delegates/gpu/cl/inference_context.cc +++ b/tensorflow/lite/delegates/gpu/cl/inference_context.cc @@ -53,22 +53,21 @@ bool IsReady(const std::unordered_set& ready_tensors, std::vector> GetCLNodeTensors( const CLNode& node) { std::vector> result; - for (int i = 0; i < node.operations.size(); ++i) { - const OperationDef op_def = node.operations[i]->GetDefinition(); - const auto& first_range = node.ranges[0]; - for (int k = first_range.x; k < first_range.y; ++k) { - result.push_back({node.inputs[k], op_def.src_tensors[k - first_range.x]}); - } - for (int j = 1; j < node.ranges.size(); ++j) { - const auto& range = node.ranges[j]; - for (int k = range.x; k < range.y; ++k) { - result.push_back({node.inputs[k], op_def.src_tensors[k - range.x + 1]}); - } - } - for (int j = 0; j < node.outputs.size(); ++j) { - result.push_back({node.outputs[j], op_def.dst_tensors[j]}); + const OperationDef main_def = node.operations[0]->GetDefinition(); + const auto& first_range = node.ranges[0]; + for (int k = first_range.x; k < first_range.y; ++k) { + result.push_back({node.inputs[k], main_def.src_tensors[k - first_range.x]}); + } + for (int j = 1; j < node.ranges.size(); ++j) { + const auto& range = node.ranges[j]; + const OperationDef op_def = node.operations[j]->GetDefinition(); + for (int k = range.x; k < range.y; ++k) { + result.push_back({node.inputs[k], op_def.src_tensors[k - range.x + 1]}); } } + for (int j = 0; j < node.outputs.size(); ++j) { + result.push_back({node.outputs[j], main_def.dst_tensors[j]}); + } return result; }