readme

README.md

@@ -77,8 +77,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model_name = "baidu/ERNIE-4.5-21B-A3B-PT"
 
 # load the tokenizer and the model
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+    tp_plan="auto",
+)
 
 # prepare the model input
 prompt = "Give me a short introduction to large language model."