Update generate.py

generate.py

@@ -9,7 +9,7 @@ def generate(model, input_ids, generation_config, **kwargs):
         logits = model(input_ids).logits
         next_token_logits = logits[:, -1, :]
         next_tokens = torch.argmax(next_token_logits, dim=-1)
-        model_inputs["input_ids"] = torch.cat((input_ids, next_tokens[:, None]), dim=-1)
+        input_ids = torch.cat((input_ids, next_tokens[:, None]), dim=-1)
         cur_length += 1
 
     return input_ids