diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py
index fdad289..a6da424 100644
--- a/llama_cpp/llama_cpp.py
+++ b/llama_cpp/llama_cpp.py
@@ -187,7 +187,7 @@ _lib.llama_get_embeddings.restype = POINTER(c_float)
 
 
 # Token Id -> String. Uses the vocabulary in the provided context
-def llama_token_to_str(ctx: llama_context_p, token: int) -> bytes:
+def llama_token_to_str(ctx: llama_context_p, token: llama_token) -> bytes:
     return _lib.llama_token_to_str(ctx, token)