et_generator: Resolved pyre errors

mlcommons · Feb 6, 2024 · 5d315b6 · 5d315b6
1 parent 2730344
commit 5d315b6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/utils/et_generator/et_generator.py b/utils/et_generator/et_generator.py
@@ -4,6 +4,7 @@
 
 from chakra.third_party.utils.protolib import encodeMessage as encode_message
 from chakra.et_def.et_def_pb2 import (
+    NodeType as ChakraNodeType,
     Node as ChakraNode,
     DoubleList,
     FloatList,
@@ -35,7 +36,7 @@
 
 NODE_ID = 0
 
-def get_node(node_name: str, node_type: int) -> ChakraNode:
+def get_node(node_name: str, node_type: ChakraNodeType) -> ChakraNode:
     global NODE_ID
     node = ChakraNode()
     node.id = NODE_ID