Update generate.py

generate.py

@@ -56,7 +56,7 @@ def custom_generate(
         generated_token_ids = torch.full((batch_size, max_new_tokens), self.tokenizer.pad_token_id, dtype=torch.long, device=device)
 
         cur_token_idx = 0
-        while cur_token_idx < max_length:
+        while cur_token_idx < max_new_tokens:
             # Sample the next token
             new_ids = self(
                 input_ids[~finished_generating],
@@ -88,9 +88,11 @@ def custom_generate(
                 if attention_mask is not None:
                     attention_mask[answer_idx, last_token_idx + 1] = 1
                 input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
-                generated_token_ids[answer_idx, cur_token_idx] = new_ids_sampled
 
-                if new_ids_sampled == self.tokenizer.eos_token_id or cur_token_idx + 1 == max_length:
+                if cur_token_idx < max_new_tokens:
+                    generated_token_ids[answer_idx, cur_token_idx] = new_ids_sampled
+
+                if new_ids_sampled == self.tokenizer.eos_token_id or cur_token_idx + 1 == max_new_tokens:
                     finished_generating[answer_idx] = 1
 
             cur_token_idx += 1