baidu
/

ERNIE-4.5-21B-A3B-Base-PT

Text Generation

Model card Files Files and versions

AntonV HF Staff commited on Jul 21

Commit

d66ca1a

·

1 Parent(s): e30c284

readme

Files changed (1) hide show

README.md +6 -1

README.md CHANGED Viewed

@@ -82,7 +82,12 @@ model_name = "baidu/ERNIE-4.5-21B-A3B-Base-PT"
 # load the tokenizer and the model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
 prompt = "Large language model is"
 model_inputs = tokenizer([prompt], add_special_tokens=False, return_tensors="pt").to(model.device)

 # 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,
+    tp_plan="auto",
+)
 prompt = "Large language model is"
 model_inputs = tokenizer([prompt], add_special_tokens=False, return_tensors="pt").to(model.device)