baidu
/

ERNIE-4.5-0.3B-PT

Text Generation

Model card Files Files and versions

AntonV HF Staff commited on Jul 21

Commit

0d4f675

·

1 Parent(s): 739a458

readme

Files changed (1) hide show

README.md +5 -1

README.md CHANGED Viewed

@@ -74,7 +74,11 @@ model_name = "baidu/ERNIE-4.5-0.3B-PT"
 # load the tokenizer and the model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
 # prepare the model input
 prompt = "Give me a short introduction to large language model."

 # load the tokenizer and the model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
 # prepare the model input
 prompt = "Give me a short introduction to large language model."