Update handler.py

handler.py

@@ -89,18 +89,17 @@ class EndpointHandler:
             prompt = self.processor.apply_chat_template(
                 messages, add_generation_prompt=True
             )
+            processed_inputs = self.processor(
+                text=prompt, images=images, return_tensors="pt"
+            ).to(self.model.device)
 
-            with torch.no_grad(), torch.autocast("cuda"):
-                processed_inputs = self.processor(
-                    text=prompt, images=images, return_tensors="pt"
-                )
-                generated_ids = self.model.generate(
-                    **processed_inputs, generation_config=generation_config
-                )
-                generated_texts = self.processor.batch_decode(
-                    generated_ids,
-                    skip_special_tokens=True,
-                )
+            generated_ids = self.model.generate(
+                **processed_inputs, generation_config=generation_config
+            )
+            generated_texts = self.processor.batch_decode(
+                generated_ids,
+                skip_special_tokens=True,
+            )
             predictions.append(generated_texts[0])
 
         return {"predictions": predictions}