diff --git a/server/text_generation_server/utils/tokens.py b/server/text_generation_server/utils/tokens.py
index 8f2ae4f6..896d4d40 100644
--- a/server/text_generation_server/utils/tokens.py
+++ b/server/text_generation_server/utils/tokens.py
@@ -232,7 +232,7 @@ class HeterogeneousNextTokenChooser:
         seeds: List[int],
         tokenizer: PreTrainedTokenizerBase,
         grammars: List[str],
-        grammar_types: List[GrammarType],
+        grammar_types: List[int],
         fsm_grammar_states=List[int],
     ):
         warpers = []