github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

add pad

pull/5190/head
Xuanlei Zhao 2023-12-25 17:02:05 +08:00
parent 23341687ed
commit b8fadb68a7
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
applications/ColossalMoE/infer.py
View File

@ -146,7 +146,8 @@ def main():
coordinator.print_on_master(f"Finish load ckpt")
text = ["Hello my name is"]
inputs = tokenizer(text, return_tensors="pt").to(torch.cuda.current_device())
tokenizer.pad_token = tokenizer.unk_token
inputs = tokenizer(text, return_tensors="pt", padding=True).to(torch.cuda.current_device())
outputs = model.module.generate(**inputs, max_new_tokens=20)
outputs = tokenizer.batch_decode(outputs)[0]
print(outputs)