alvarobartt HF staff commited on
Commit
89baa6d
·
verified ·
1 Parent(s): 64f317d

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -31,7 +31,7 @@ class EndpointHandler:
31
  self.model = AutoModelForVision2Seq.from_pretrained(
32
  model_dir,
33
  torch_dtype=torch.bfloat16,
34
- _attn_implementation="flash_attention_2",
35
  device_map="auto",
36
  ).eval()
37
  self.generation_config = GenerationConfig.from_pretrained(model_dir)
 
31
  self.model = AutoModelForVision2Seq.from_pretrained(
32
  model_dir,
33
  torch_dtype=torch.bfloat16,
34
+ _attn_implementation="eager", # "flash_attention_2",
35
  device_map="auto",
36
  ).eval()
37
  self.generation_config = GenerationConfig.from_pretrained(model_dir)