Update handler.py

handler.py

@@ -27,8 +27,11 @@ class EndpointHandler():
 
         # Load the Qwen2-VL model
         self.model = Qwen2VLForConditionalGeneration.from_pretrained(
-            self.model_dir, torch_dtype="auto", device_map="auto"
-        )
+                 self.model_dir,
+                 torch_dtype=torch.bfloat16,
+                 attn_implementation="flash_attention_2",
+                 device_map="auto",
+         )
         self.processor = AutoProcessor.from_pretrained(self.model_dir)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: