[BugFix][Runtime] Fix Incorrect node information

apache · Jan 4, 2023 · 440e2ed · 440e2ed
1 parent 49ed544
commit 440e2ed
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/python/tvm/contrib/debugger/debug_result.py b/python/tvm/contrib/debugger/debug_result.py
@@ -73,21 +73,25 @@ def _update_graph_json(self):
         """update the nodes_list with name, shape and data type,
         for temporarily storing the output.
         """
-        nodes_len = len(self._nodes_list)
-        for i in range(nodes_len):
-            node = self._nodes_list[i]
+        eid = 0
+        for node in self._nodes_list:
             input_list = []
-            for input_node in node["inputs"]:
-                input_list.append(self._nodes_list[input_node[0]]["name"])
-            node["inputs"] = input_list
-            dtype = str("type: " + self._dtype_list[1][i])
-            if "attrs" not in node:
+            if node["op"] == "null":
                 node["attrs"] = {}
                 node["op"] = "param"
-            else:
+                num_outputs = 1
+            elif node["op"] == "tvm_op":
+                for input_node in node["inputs"]:
+                    input_list.append(self._nodes_list[input_node[0]]["name"])
                 node["op"] = node["attrs"]["func_name"]
+                num_outputs = int(node["attrs"]["num_outputs"])
+            else:
+                raise ValueError("")
+            node["inputs"] = input_list
+            dtype = str("type: " + self._dtype_list[1][eid])
             node["attrs"].update({"T": dtype})
-            node["shape"] = self._shapes_list[1][i]
+            node["shape"] = self._shapes_list[1][eid]
+            eid += num_outputs
 
     def _cleanup_tensors(self):
         """Remove the tensor dump file (graph wont be removed)"""