HongxuanLi
/

nougat-base-deploy

vision-encoder-decoder

image-text-to-text

Model card Files Files and versions

Hongxuan Li commited on Apr 22, 2024

Commit

9308c9e

·

1 Parent(s): d21a961

add handler

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -10,7 +10,7 @@ class EndpointHandler():
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.processor = NougatProcessor.from_pretrained(path)
         self.model = VisionEncoderDecoderModel.from_pretrained(path)
-        self.model = model.to(self.device)
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
         Args:

         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.processor = NougatProcessor.from_pretrained(path)
         self.model = VisionEncoderDecoderModel.from_pretrained(path)
+        self.model = self.model.to(self.device)
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
         Args: