Add bf16

README.md

@@ -112,12 +112,12 @@ export use_triton=False
 ```python
 >>> from transformers import AutoModelForCausalLM, AutoTokenizer
 >>> tokenizer = AutoTokenizer.from_pretrained("OpenNLPLab/TransNormerLLM2-7B-300B", trust_remote_code=True)
->>> model = AutoModelForCausalLM.from_pretrained("TransNormerLLM2-7B-300B", device_map="auto", trust_remote_code=True)
+>>> model = AutoModelForCausalLM.from_pretrained("TransNormerLLM2-7B-300B", torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
 >>> inputs = tokenizer('今天是美好的一天', return_tensors='pt')
 >>> pred = model.generate(**inputs, max_new_tokens=8192, repetition_penalty=1.0)
 >>> print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
 ```
-
+* **Note**: we recommend to use `bfloat16` in `TransNormerLLM`, `float16` might lead `nan` error, please check your divce compatibility!
 
 # Fine-tuning the Model