prithivMLmods commited on
Commit
24147c0
·
verified ·
1 Parent(s): 7288593

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -0
app.py CHANGED
@@ -121,6 +121,7 @@ MODEL_PATH = 'tencent/POINTS-Reader'
121
  print(f"Loading model: {MODEL_PATH}")
122
  model = AutoModelForCausalLM.from_pretrained(
123
  MODEL_PATH,
 
124
  trust_remote_code=True,
125
  torch_dtype=torch.float16,
126
  device_map='auto'
 
121
  print(f"Loading model: {MODEL_PATH}")
122
  model = AutoModelForCausalLM.from_pretrained(
123
  MODEL_PATH,
124
+ attn_implementation="flash_attention_2",
125
  trust_remote_code=True,
126
  torch_dtype=torch.float16,
127
  device_map='auto'