Update handler.py
Browse files- 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)
|