Gong Baitao
commited on
Commit
•
b8e0001
1
Parent(s):
94ea18d
Update README.md
Browse files
README.md
CHANGED
@@ -45,8 +45,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
45 |
from accelerate import dispatch_model
|
46 |
from accelerate.utils import get_balanced_memory, infer_auto_device_map
|
47 |
|
48 |
-
tokenizer = AutoTokenizer.from_pretrained("openbmb/cpm-bee-
|
49 |
-
model = AutoModelForCausalLM.from_pretrained("openbmb/cpm-bee-
|
50 |
|
51 |
max_memory = get_balanced_memory(
|
52 |
model,
|
|
|
45 |
from accelerate import dispatch_model
|
46 |
from accelerate.utils import get_balanced_memory, infer_auto_device_map
|
47 |
|
48 |
+
tokenizer = AutoTokenizer.from_pretrained("openbmb/cpm-bee-10b", trust_remote_code=True)
|
49 |
+
model = AutoModelForCausalLM.from_pretrained("openbmb/cpm-bee-10b", trust_remote_code=True).cuda()
|
50 |
|
51 |
max_memory = get_balanced_memory(
|
52 |
model,
|