[BUG] 我下载了huggingface上的baichuan7b模型，使用 里面的测试程序测试发现CUDA错误 #143
Description
Required prerequisites
- I have read the documentation https://github.com/baichuan-inc/baichuan-7B/blob/HEAD/README.md.
- I have searched the Issue Tracker and Discussions that this hasn't already been reported. (+1 or comment there if it has.)
- Consider asking first in a Discussion.
System information
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
tokenizer = AutoTokenizer.from_pretrained("./", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("./", device_map="auto", trust_remote_code=True)
inputs = tokenizer('登鹳雀楼->王之涣\n夜雨寄北->', return_tensors='pt')
inputs = inputs.to('cuda:2')
pred = model.generate(**inputs, max_new_tokens=16,repetition_penalty=1.1)
print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
Problem description
Reproducible example code
The Python snippets:
Command lines:
Extra dependencies:
Steps to reproduce:
Traceback
No response
Expected behavior
No response
Additional context
No response
Checklist
- I have provided all relevant and necessary information above.
- I have chosen a suitable title for this issue.