model to gpu

handler.py

@@ -1,4 +1,5 @@
 import os
+import torch
 from typing import Dict, List, Any
 from transformers import AutoTokenizer
 from gector import GECToR, predict, load_verb_dict
@@ -12,6 +13,9 @@ class EndpointHandler:
             os.path.join(path, "data/verb-form-vocab.txt")
         )
 
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(device)
+
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Process the input data and return the predicted results.