diff --git a/hta/common/execution_trace.py b/hta/common/execution_trace.py
index 17ab9b2..1e9f4d2 100644
--- a/hta/common/execution_trace.py
+++ b/hta/common/execution_trace.py
@@ -162,7 +162,7 @@ def add_et_column(trace_df: pd.DataFrame, et: ExecutionTrace, column: str) -> No
         return
     if column == "op_schema":
 
-        def map_func(node_id):
+        def map_func(node_id) -> str:
             return et.nodes[node_id].op_schema
 
     elif column == "input_shapes":
@@ -182,7 +182,7 @@ def map_func(node_id):
 
     elif column == "output_types":
 
-        def map_func(node_id):
+        def map_func(node_id) -> List[str]:
             return et.nodes[node_id].output_types
 
     elif column == "et_node_name":