Add torch_dtype to README example

README.md

@@ -75,13 +75,15 @@ Overall, InternLM-20B comprehensively outperforms open-source models in the 13B
 ## Import from Transformers
 To load the InternLM 20B model using Transformers, use the following code:
 ```python
->>> from transformers import AutoTokenizer, AutoModelForCausalLM
->>> tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-chat-20b", trust_remote_code=True)
->>> model = AutoModelForCausalLM.from_pretrained("internlm/internlm-chat-20b", trust_remote_code=True).cuda()
->>> model = model.eval()
->>> output, history = model.chat(tokenizer, "Hello! Today is sunny, it is time to go out")
->>> print(output)
-Hello! Today is sunny, and it sounds like a great day to go out an enjoy the weather. What would you like to do?
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-chat-20b", trust_remote_code=True)
+# Set `torch_dtype=torch.bfloat16` to load model in bfloat16, otherwise it will be loaded as float32 and cause OOM Error.
+model = AutoModelForCausalLM.from_pretrained("internlm/internlm-chat-20b", torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
+model = model.eval()
+output, history = model.chat(tokenizer, "Hello! Today is sunny, it is time to go out")
+print(output)
+# Hello! Today is sunny, and it sounds like a great day to go out an enjoy the weather. What would you like to do?
 ```
 
 **Limitations:** Although we have made efforts to ensure the safety of the model during the training process and to encourage the model to generate text that complies with ethical and legal requirements, the model may still produce unexpected outputs due to its size and probabilistic generation paradigm. For example, the generated responses may contain biases, discrimination, or other harmful content. Please do not propagate such content. We are not responsible for any consequences resulting from the dissemination of harmful information.
@@ -140,13 +142,15 @@ InternLM 20B 在模型结构上选择了深结构，层数设定为60层，超
 ## 通过 Transformers 加载
 通过以下的代码加载 InternLM 20B 模型
 ```python
->>> from transformers import AutoTokenizer, AutoModelForCausalLM
->>> tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-chat-20b", trust_remote_code=True)
->>> model = AutoModelForCausalLM.from_pretrained("internlm/internlm-chat-20b", trust_remote_code=True).cuda()
->>> model = model.eval()
->>> output, history = model.chat(tokenizer, "你好呀！今天天气真好")
->>> print(output)
-你好！是的，今天的天气非常晴朗，非常适合户外活动。
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-chat-20b", trust_remote_code=True)
+# `torch_dtype=torch.bfloat16` 可以令模型以 bfloat16 精度加载，否则 transformers 会将模型加载为 float32，导致显存不足
+model = AutoModelForCausalLM.from_pretrained("internlm/internlm-chat-20b", torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
+model = model.eval()
+output, history = model.chat(tokenizer, "你好呀！今天天气真好")
+print(output)
+# 你好！是的，今天的天气非常晴朗，非常适合户外活动。
 ```
 
 **局限性：** 尽管在训练过程中我们非常注重模型的安全性，尽力促使模型输出符合伦理和法律要求的文本，但受限于模型大小以及概率生成范式，模型可能会产生各种不符合预期的输出，例如回复内容包含偏见、歧视等有害内容，请勿传播这些内容。由于传播不良信息导致的任何后果，本项目不承担责任。