"."

model.py

@@ -164,8 +164,8 @@ class GPT(PreTrainedModel):
         return {"input_ids": input_ids, "past_key_values": past_key_values}
 
     @torch.no_grad()
-    def generate(self, input_ids, max_new_tokens, temperature=1.0, top_k=None, attention_mask=None):
-        for _ in range(max_new_tokens):
+    def generate(self, input_ids, max_length, temperature=1.0, top_k=None, attention_mask=None):
+        for _ in range(max_length - input_ids.size(1)):
             idx_cond = input_ids if input_ids.size(1) <= self.config.block_size else input_ids[:, -self.config.block_size:]
             out = self(idx_cond)
             logits = out['logits'][:, -1, :] / temperature