Move self.predictions above the make_graph

zhiqwang · Sep 16, 2021 · 31698fc · 31698fc
1 parent 8bd2673
commit 31698fc
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/yolort/relaying/ir_visualizer.py b/yolort/relaying/ir_visualizer.py
@@ -93,6 +93,10 @@ def make_graph(self, module, dot=None, parent_dot=None, prefix="", input_preds=N
 
                     # go into subgraph
                     sub_prefix = f'{prefix}{submodule_name}.'
+
+                    for i, o in enumerate(node.outputs()):
+                        self.predictions[o] = {f'{sub_prefix}output_{i}'}, set()
+
                     with dot.subgraph(name=f'cluster_{name}') as sub_dot:
                         sub_dot.attr(label=label)
                         sub_module = module
@@ -109,8 +113,6 @@ def make_graph(self, module, dot=None, parent_dot=None, prefix="", input_preds=N
                             classes_found=classes_found,
                         )
 
-                    for i, o in enumerate(node.outputs()):
-                        self.predictions[o] = {f'{sub_prefix}output_{i}'}, set()
                 else:
                     dot.node(name, label=label, shape='box')
                     for i in relevant_inputs: