Simplify TF get_output_names (apache#3025)

neo-ai · May 13, 2019 · c9798ad · c9798ad
1 parent f9518dc
commit c9798ad
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/python/tvm/relay/frontend/tensorflow_parser.py b/python/tvm/relay/frontend/tensorflow_parser.py
@@ -85,11 +85,9 @@ def _get_output_names(self):
                                                         tags,
                                                         self._model_dir)
             output_names = set()
-            for k in meta_graph_def.signature_def.keys():
-                outputs_tensor_info = meta_graph_def.signature_def[k].outputs
-                for output_tensor in outputs_tensor_info.values():
-                    output_names.add(output_tensor.name)
-            output_names = [i.replace(":0", "") for i in output_names]
+            for sig_def in meta_graph_def.signature_def.values():
+                for output_tensor in sig_def.outputs.values():
+                    output_names.add(output_tensor.name.replace(":0", ""))
             return ",".join(output_names)
 
     def _load_saved_model(self):