Update generate.py

generate.py

@@ -47,6 +47,8 @@ def custom_generate(
     with torch.no_grad():
         finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
 
+        if max_new_tokens is None:
+            max_new_tokens = 50  # Default value if not specified
         for cur_token_idx in range(max_new_tokens):
             # Sample the next token
             new_ids = self(