diff --git a/jsonformer/main.py b/jsonformer/main.py
index 9c13471..25fee0b 100644
--- a/jsonformer/main.py
+++ b/jsonformer/main.py
@@ -59,7 +59,7 @@ def generate_number(self, temperature: Union[float, None] = None, iterations=0):
             self.model.device
         )
         response = self.model.generate(
-            input_tokens,
+            inputs=input_tokens,
             max_new_tokens=self.max_number_tokens,
             num_return_sequences=1,
             logits_processor=[self.number_logit_processor],
@@ -110,7 +110,7 @@ def generate_string(self) -> str:
         )
 
         response = self.model.generate(
-            input_tokens,
+            inputs=input_tokens,
             max_new_tokens=self.max_string_token_length,
             num_return_sequences=1,
             temperature=self.temperature,