jasper-lu
/

pix2struct_embedding

Model card Files Files and versions

Jasper Lu commited on Nov 15, 2023

Commit

ac64690

·

1 Parent(s): 4efd111

to cuda

Files changed (1) hide show

handler.py +2 -1

handler.py CHANGED Viewed

@@ -17,8 +17,9 @@ class EndpointHandler():
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         url = data.pop("inputs", data)
         image = Image.open(requests.get(url, stream=True).raw)
-        inputs = self.processor(images=image, return_tensors="pt").cuda()
         with torch.no_grad():
             outputs = self.model(**inputs)

     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         url = data.pop("inputs", data)
+        device = "cuda"
         image = Image.open(requests.get(url, stream=True).raw)
+        inputs = self.processor(images=image, return_tensors="pt").to(device)
         with torch.no_grad():
             outputs = self.model(**inputs)