Update README.md

README.md

@@ -16,8 +16,10 @@ Then simply use the following code to load the model:
 ```python
 from modeling_memoryllm import MemoryLLM
 from transformers import AutoTokenizer
-model = MemoryLLM.from_pretrained("YuWangX/memoryllm-8b-chat")
+# load chat model
+model = MemoryLLM.from_pretrained("YuWangX/memoryllm-8b-chat", attn_implementation="flash_attention_2", torch_dtype=torch.float16)
 tokenizer = AutoTokenizer.from_pretrained("YuWangX/memoryllm-8b-chat")
+model = model.cuda()
 ```
 
 ### How to use the model