Update generate.py

generate.py

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