From 30c25bee177e183ed0854a2784805c0f31c704e8 Mon Sep 17 00:00:00 2001 From: Julian Date: Mon, 16 Nov 2020 09:59:26 +0100 Subject: [PATCH] handle list of names --- keras2onnx/_parser_tf.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/keras2onnx/_parser_tf.py b/keras2onnx/_parser_tf.py index 38b686cb..153ae6e8 100644 --- a/keras2onnx/_parser_tf.py +++ b/keras2onnx/_parser_tf.py @@ -289,10 +289,12 @@ def extract_outputs_from_inbound_nodes(model): output_tensors = [nd_.output_tensors] if hasattr(nd_.output_tensors, 'dtype') else \ nd_.output_tensors for ts_ in output_tensors: - op_name = tsname_to_node(ts_.name) - if op_name not in output_dict: - output_dict[op_name] = (model, None) - + if not isinstance(ts_, (tuple, list)): + ts_ = [ts_] + for ts__ in ts_: + op_name = tsname_to_node(ts__.name) + if op_name not in output_dict: + output_dict[op_name] = (model, None) for ts_ in _get_output_mask(model): if ts_ is not None: output_dict[ts_.op.name] = (model, model)