Another GPU test

handler.py

@@ -19,12 +19,12 @@ class EndpointHandler():
         image_input = inputs["image"] if "image" in inputs else None
 
         if text_input:
-            processor = self.processor(text=text_input, return_tensors="pt", padding=True)
+            processor = self.processor(text=text_input, return_tensors="pt", padding=True).to(device)
             with torch.no_grad():
                 return self.text_model(**processor).pooler_output.tolist()
         elif image_input:
             image = Image.open(BytesIO(base64.b64decode(image_input)))
-            processor = self.processor(images=image, return_tensors="pt")
+            processor = self.processor(images=image, return_tensors="pt").to(device)
             with torch.no_grad():
                 return self.image_model(**processor).image_embeds.tolist()
         else: