Update README.md for flash attn

README.md

@@ -105,10 +105,16 @@ triton==2.0.0.dev20221202
 
 Then, move the model to `bfloat16` and use it as follows:
 ```python
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoConfig
+
+config = AutoConfig.from_pretrained(
+    "replit/replit-code-v1-3b",
+    trust_remote_code=True
+)
+config.attn_config['attn_impl'] = 'triton'
 
 # load model
-model = AutoModelForCausalLM.from_pretrained('replit/replit-code-v1-3b', trust_remote_code=True, attn_impl='triton')
+model = AutoModelForCausalLM.from_pretrained('replit/replit-code-v1-3b', config=config, trust_remote_code=True)
 model.to(device='cuda:0', dtype=torch.bfloat16)
 
 # forward pass