change max_length logic

custom_generate/generate.py

@@ -44,7 +44,7 @@ def generate(model, input_ids, generation_config=None, n_sigma:float=1.0, **kwar
     """
     generation_config = generation_config or model.generation_config  # default to the model generation config
     cur_length = input_ids.shape[1]
-    if generation_config.max_length is None:
+    if generation_config.max_new_tokens:
         max_length = cur_length + generation_config.max_new_tokens
     else:
         max_length = generation_config.max_length