[update]edit main

main.py

@@ -78,6 +78,9 @@ def main():
             return_tensors="pt",
             add_special_tokens=False,
         ).input_ids.to(args.device)
+        bos_token_id = torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long).to(args.device)
+        eos_token_id = torch.tensor([[tokenizer.eos_token_id]], dtype=torch.long).to(args.device)
+        input_ids = torch.concat([bos_token_id, input_ids, eos_token_id], dim=1)
 
         with torch.no_grad():
             outputs = model.generate(