Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -39,6 +39,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
39 |
model = AutoModelForCausalLM.from_pretrained(
|
40 |
model_id,
|
41 |
device_map="auto",
|
|
|
42 |
attn_implementation="flash_attention_2",
|
43 |
# quantization_config=quantization_config,
|
44 |
)
|
|
|
39 |
model = AutoModelForCausalLM.from_pretrained(
|
40 |
model_id,
|
41 |
device_map="auto",
|
42 |
+
torch_dtype=torch.bfloat16,
|
43 |
attn_implementation="flash_attention_2",
|
44 |
# quantization_config=quantization_config,
|
45 |
)
|