Update handler.py

handler.py

@@ -9,7 +9,11 @@ class EndpointHandler():
         model = AutoModelForMaskedLM.from_pretrained(path)
         self.tokenizer = tokenizer
         self.model = model
-
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda")
+            self.model.to(self.device)
+        else:
+            self.device = torch.device("cpu")
     def __call__(self, data: Dict[str, Any]) -> List[Dict[Any, Any]]:
         """
        data args:
@@ -20,7 +24,7 @@ class EndpointHandler():
         """
         # get inputs
         text = data.pop("text", data)
-        tokens = self.tokenizer(text, return_tensors='pt', padding=True)
+        tokens = self.tokenizer(text, return_tensors='pt', padding=True).to(self.device)
         outputs = self.model(**tokens)
         results = []
         for idx, x in enumerate(outputs.logits):