tangzhy commited on
Commit
5ba4c95
·
verified ·
1 Parent(s): a9764a0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -0
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
  )