added device to handler

handler.py

@@ -5,8 +5,12 @@ import torch
 
 class EndpointHandler:
     def __init__(self, path=""):
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+        self.tokenizer = AutoTokenizer.from_pretrained(path, device=device)
         self.model = AutoModelForSequenceClassification.from_pretrained(path)
+        self.model.to(device)
+
         def tokenize(batch):
             return self.tokenizer(
                 batch['topic'],
@@ -16,7 +20,8 @@ class EndpointHandler:
                 return_offsets_mapping=False,
                 padding="max_length",
                 return_tensors='pt'
-            )
+            ).to(device)
+        
         self.tokenize = tokenize
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
@@ -36,11 +41,11 @@ class EndpointHandler:
         }
 
         for topic in topics:
-          for text in texts:        
+          for text in texts:
             batch['id'].append(text['id'])
             batch['text'].append(text['text'])
             batch['topic'].append(topic)
-        
+
         tokenized_inputs = self.tokenize(batch)
 
         # run normal prediction
@@ -48,4 +53,4 @@ class EndpointHandler:
         predictions = torch.argmax(output.logits, dim=-1).numpy(force=True)
         batch['label'] = [self.model.config.id2label[p] for p in predictions]
         batch.pop('text')
-        return batch
+        return batch