manu commited on
Commit
19e6f48
·
verified ·
1 Parent(s): 89cecf3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -0
app.py CHANGED
@@ -17,6 +17,7 @@ model = ColQwen2.from_pretrained(
17
  "manu/colqwen2-v1.0-alpha",
18
  torch_dtype=torch.bfloat16,
19
  device_map="cuda:0", # or "mps" if on Apple Silicon
 
20
  ).eval()
21
  processor = ColQwen2Processor.from_pretrained("manu/colqwen2-v1.0-alpha")
22
 
 
17
  "manu/colqwen2-v1.0-alpha",
18
  torch_dtype=torch.bfloat16,
19
  device_map="cuda:0", # or "mps" if on Apple Silicon
20
+ attn_implementation="flash_attention_2", # should work on A100
21
  ).eval()
22
  processor = ColQwen2Processor.from_pretrained("manu/colqwen2-v1.0-alpha")
23