Update handler.py

handler.py

@@ -18,7 +18,7 @@ class EndpointHandler():
         if "image" in inputs:
             # decode base64 image to PIL
             image = Image.open(BytesIO(base64.b64decode(inputs['image'])))  
-            inputs = self.processor(images=image, text=None, return_tensors="pt", padding=True)
+            inputs = self.processor(images=image, text=None, return_tensors="pt", padding=True).to(device)
             
             image_embeds = self.model.get_image_features(
                 pixel_values=inputs["pixel_values"]
@@ -27,7 +27,7 @@ class EndpointHandler():
             return image_embeds[0].tolist()
         if "text" in inputs:
             text = inputs['text']
-            inputs = self.processor(images=None, text=text, return_tensors="pt", padding=True)
+            inputs = self.processor(images=None, text=text, return_tensors="pt", padding=True).to(device)
             
             text_embeds = self.model.get_text_features(
                 input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"]