Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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="
|
| 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
|