diff --git a/template/pipelines/inference.py b/template/pipelines/inference.py
index 691f58b..3449a2c 100644
--- a/template/pipelines/inference.py
+++ b/template/pipelines/inference.py
@@ -12,7 +12,7 @@
 
 
 @pipeline
-def inference(random_state: str, target: str):
+def inference(random_state: int, target: str):
     """
     Model inference pipeline.