Update tamilatis/predict.py

tamilatis/predict.py

@@ -16,10 +16,10 @@ class TamilATISPredictor:
         num_labels,
     ):
         self.model = model
-        self.model.load_state_dict(torch.load(checkpoint_path))
         self.model.eval()
         self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = "cuda" if torch.cuda.is_available() else "CPU"
+        self.model.load_state_dict(torch.load(checkpoint_path,map_location=self.device))
         self.num_labels = num_labels
         self.label_encoder = label_encoder
         self.intent_encoder = intent_encoder