m-ric commited on
Commit
e6c50ab
·
verified ·
1 Parent(s): 38591ef

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -13,11 +13,12 @@ import spaces
13
  processor = AutoProcessor.from_pretrained(
14
  "rhymes-ai/Aria",
15
  revision="0e2a8319",
 
16
  )
17
  model = AutoModel.from_pretrained("rhymes-ai/Aria",
18
  revision="f7c92bba",
19
  torch_dtype=torch.bfloat16,
20
- #_attn_implementation="flash_attention_2"
21
  ).to("cuda")
22
 
23
  @spaces.GPU
 
13
  processor = AutoProcessor.from_pretrained(
14
  "rhymes-ai/Aria",
15
  revision="0e2a8319",
16
+ torch_dtype=torch.bfloat16,
17
  )
18
  model = AutoModel.from_pretrained("rhymes-ai/Aria",
19
  revision="f7c92bba",
20
  torch_dtype=torch.bfloat16,
21
+ attn_implementation="flash_attention_2"
22
  ).to("cuda")
23
 
24
  @spaces.GPU