Get pg_name from record_param_comms for collectives

mlcommons · Nov 16, 2024 · 396bc6e · 396bc6e
1 parent 65e2c00
commit 396bc6e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/converter/pytorch_converter.py b/src/converter/pytorch_converter.py
@@ -346,6 +346,8 @@ def get_protobuf_node_type_from_json_node(
                     if parent_node.name == "record_param_comms"
                     else parent_node.name
                 )
+                if parent_node.name == "record_param_comms" and parent_node.pg_name is not "":
+                    json_node.pg_name = parent_node.pg_name
                 if "send" in keyword:
                     return COMM_SEND_NODE
                 if "recv" in keyword: