teowu commited on
Commit
9dd6788
1 Parent(s): edd457f

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -1
README.md CHANGED
@@ -11,7 +11,8 @@ from transformers import AutoModelForCausalLM
11
 
12
  model = AutoModelForCausalLM.from_pretrained("q-future/co-instruct-preview",
13
  trust_remote_code=True,
14
- torch_dtype=torch.float16,
 
15
  device_map={"":"cuda:0"})
16
  ```
17
 
 
11
 
12
  model = AutoModelForCausalLM.from_pretrained("q-future/co-instruct-preview",
13
  trust_remote_code=True,
14
+ torch_dtype=torch.float16,
15
+ attn_implementation="flash_attention_2",
16
  device_map={"":"cuda:0"})
17
  ```
18