diff --git a/server/text_generation_server/models/vectorized_causal_lm.py b/server/text_generation_server/models/vectorized_causal_lm.py
index 5abc066c..2c8477a9 100644
--- a/server/text_generation_server/models/vectorized_causal_lm.py
+++ b/server/text_generation_server/models/vectorized_causal_lm.py
@@ -117,7 +117,7 @@ class VectorizedCausalLMBatch(Batch):
             token_offsets=token_offsets,
             next_token_chooser=next_token_chooser,
             stopping_criterias=stopping_criterias,
-            max_input_length=max_input_length.item(),
+            max_input_length=max_input_length,
             max_tokens=max_tokens,
         )
 
@@ -483,8 +483,6 @@ class VectorizedNextTokenChooser:
         )
 
 
-
-
 class VectorizedCausalLM(Model):
     def __init__(
         self,