Merge pull request #16 from JoongunPark/main

[chakra][et_converter] Add Annotation for Graph Trace Observer in PyTorch
mlcommons · Jan 10, 2024 · 8f0edff · 8f0edff
2 parents e7d7a7c + e3f7039
commit 8f0edff
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/et_converter/pytorch2chakra_converter.py b/et_converter/pytorch2chakra_converter.py
@@ -382,6 +382,8 @@ def find_root_nids(
         for node in nodes:
             if "[pytorch|profiler|execution_graph|thread]" in node["name"]:
                 root_nids.append(node["id"])
+            elif "[pytorch|profiler|execution_trace|thread]" in node["name"]:
+                root_nids.append(node["id"])
         if not root_nids:
             raise ValueError("Cannot find a root NID")
         return root_nids