Fix for the test

custom_generate/generate.py

@@ -6,7 +6,7 @@ def generate(model, input_ids, generation_config=None, left_padding=None, **kwar
     
     generation_config = generation_config or model.generation_config  # default to the model generation config
     cur_length = input_ids.shape[1]
-    max_length = generation_config.max_length or cur_length + generation_config.max_new_tokens
+    max_length = generation_config.max_length or cur_length + kwargs.get("max_new_tokens")
 
     # Example of custom argument: add `left_padding` (integer) pad tokens before the prompt
     if left_padding is not None: