RYDE-WORK/MiniCPM
1
0
Fork 0
mirror of https://github.com/RYDE-WORK/MiniCPM.git synced 2026-01-19 12:53:36 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

use default gpu 0

Browse Source
This commit is contained in:
winter 2024-07-04 19:54:18 +08:00
parent a5c9423f7e
commit a77a34759f
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
demo/hf_based_demo.py
View File

@ -34,7 +34,7 @@ else:
# init model and tokenizer
path = args.model_path
tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch_dtype, device_map="auto", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch_dtype, device_map="cuda:0", trust_remote_code=True)
model_architectures = model.config.architectures[0]