vikhyatk commited on
Commit
1e2b81a
1 Parent(s): b35f285

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -17,7 +17,8 @@ model_id = "vikhyatk/moondream2"
17
  revision = "2024-04-02"
18
  tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
19
  moondream = AutoModelForCausalLM.from_pretrained(
20
- model_id, trust_remote_code=True, revision=revision
 
21
  ).to(device=device, dtype=dtype)
22
  moondream.eval()
23
 
 
17
  revision = "2024-04-02"
18
  tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
19
  moondream = AutoModelForCausalLM.from_pretrained(
20
+ model_id, trust_remote_code=True, revision=revision,
21
+ attn_implementation="flash_attention_2"
22
  ).to(device=device, dtype=dtype)
23
  moondream.eval()
24