Add stream ID to GPU operators in the converter

mlcommons · Jul 11, 2024 · c27b8fb · c27b8fb
1 parent 7f8b892
commit c27b8fb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/converter/pytorch_converter.py b/src/converter/pytorch_converter.py
@@ -337,6 +337,8 @@ def convert_to_chakra_node(
                 ChakraAttr(name="is_cpu_op", bool_val=not pytorch_node.is_gpu_op()),
             ]
         )
+        if pytorch_node.stream is not None:
+            chakra_node.attr.append(ChakraAttr(name="stream", int64_val=pytorch_node.stream))
         return chakra_node
 
     def get_chakra_node_type_from_pytorch_node(