diff --git a/packages-nextgen/kestrel_core/src/kestrel/ir/graph.py b/packages-nextgen/kestrel_core/src/kestrel/ir/graph.py
index d555c2f6..645f7454 100644
--- a/packages-nextgen/kestrel_core/src/kestrel/ir/graph.py
+++ b/packages-nextgen/kestrel_core/src/kestrel/ir/graph.py
@@ -191,7 +191,7 @@ def get_nodes_by_type(self, ntype: type) -> Iterable[Instruction]:
         return [n for n in self.nodes() if isinstance(n, ntype)]
 
     def get_nodes_by_type_and_attributes(
-        self, ntype: type, attr2val: Mapping[str, Union[str, bool]]
+        self, ntype: type, attr2val: Mapping[str, Union[str, bool, int]]
     ) -> Iterable[Instruction]:
         """Get nodes by both type and attributes/values
 
diff --git a/packages-nextgen/kestrel_core/src/kestrel/ir/instructions.py b/packages-nextgen/kestrel_core/src/kestrel/ir/instructions.py
index e0f874e7..0f8f7b59 100644
--- a/packages-nextgen/kestrel_core/src/kestrel/ir/instructions.py
+++ b/packages-nextgen/kestrel_core/src/kestrel/ir/instructions.py
@@ -181,7 +181,7 @@ def get_instruction_class(name: str) -> Type[Instruction]:
 
 
 @typechecked
-def instruction_from_dict(d: Mapping[str, Union[str, bool]]) -> Instruction:
+def instruction_from_dict(d: Mapping[str, Union[str, bool, int]]) -> Instruction:
     instruction_class = get_instruction_class(d["instruction"])
     try:
         instruction = instruction_class.from_dict(d)