robinhad commited on
Commit
dc5393d
·
verified ·
1 Parent(s): b8addeb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -28,7 +28,7 @@ def load_model():
28
  MODEL_ID,
29
  torch_dtype=torch.bfloat16, # if device == "cuda" else torch.float32,
30
  device_map="auto", # if device == "cuda" else None,
31
- attn_implementation="fa3", # "flash_attention_2", #
32
  ) # .cuda()
33
  print(f"Selected device:", device)
34
  return model, tokenizer, device
 
28
  MODEL_ID,
29
  torch_dtype=torch.bfloat16, # if device == "cuda" else torch.float32,
30
  device_map="auto", # if device == "cuda" else None,
31
+ attn_implementation="flash_attention_3", # "flash_attention_2", #
32
  ) # .cuda()
33
  print(f"Selected device:", device)
34
  return model, tokenizer, device