Home Explore Help
Register Sign In
marx
/
llama-recipes
1
0
Fork 0
Files Issues 0 Pull Requests 0 Wiki
Browse Source

Fix the chat example for generate when load peft model (#158)

Geeta Chauhan 1 year ago
parent
a14363ca41 0e9d1dfa78
commit
322522e9a2
1 changed files with 1 additions and 1 deletions
Split View Show Diff Stats
  1. 1 1
      inference/chat_completion.py

+ 1 - 1
inference/chat_completion.py
View File 

@@ -107,7 +107,7 @@ def main(
 
             tokens= tokens.unsqueeze(0)
 
             tokens= tokens.to("cuda:0")
 
             outputs = model.generate(
 
-                tokens,
 
+                input_ids=tokens,
 
                 max_new_tokens=max_new_tokens,
 
                 do_sample=do_sample,
 
                 top_p=top_p,