Support 1.0.2-chakra.0.0.4

et_converter/pytorch2chakra_converter.py

@@ -236,7 +236,8 @@ def _parse_and_instantiate_nodes(self, pytorch_et_data: Dict) -> None:
 
         pytorch_nodes = pytorch_et_data["nodes"]
         pytorch_node_objects = {
-            node_data["id"]: PyTorchNode(node_data) for node_data in pytorch_nodes
+            node_data["id"]: PyTorchNode(self.pytorch_schema, node_data)
+            for node_data in pytorch_nodes
         }
         self._establish_parent_child_relationships(pytorch_node_objects)
 
@@ -332,7 +333,10 @@ def convert_to_chakra_node(self, pytorch_node: PyTorchNode) -> ChakraNode:
         chakra_node.type = self.get_chakra_node_type_from_pytorch_node(pytorch_node)
         if pytorch_node.parent in self.chakra_nodes:
             chakra_node.ctrl_deps.append(pytorch_node.parent)
-        chakra_node.duration_micros = pytorch_node.exclusive_dur
+        if hasattr(pytorch_node, 'exclusive_dur'):
+            chakra_node.duration_micros = pytorch_node.exclusive_dur
+        else:
+            chakra_node.duration_micros = 0
         chakra_node.inputs.values = str(pytorch_node.inputs)
         chakra_node.inputs.shapes = str(pytorch_node.input_shapes)
         chakra_node.inputs.types = str(pytorch_node.input_types)
@@ -348,8 +352,9 @@ def convert_to_chakra_node(self, pytorch_node: PyTorchNode) -> ChakraNode:
             ChakraAttr(name="fw_tid", int64_val=pytorch_node.fw_tid),
             ChakraAttr(name="op_schema", string_val=pytorch_node.op_schema),
             ChakraAttr(name="is_cpu_op", int32_val=not pytorch_node.is_gpu_op()),
-            ChakraAttr(name="ts", int64_val=pytorch_node.ts)
         ])
+        if hasattr(pytorch_node, 'ts'):
+            chakra_node.attr.append(ChakraAttr(name="ts", int64_val=pytorch_node.ts))
         return chakra_node
 
     def get_chakra_node_type_from_pytorch_node(self, pytorch_node: PyTorchNode) -> ChakraNodeType: