seanbenhur commited on
Commit
af6f767
1 Parent(s): c64606c

Update tamilatis/predict.py

Browse files
Files changed (1) hide show
  1. tamilatis/predict.py +1 -1
tamilatis/predict.py CHANGED
@@ -18,7 +18,7 @@ class TamilATISPredictor:
18
  self.model = model
19
  self.model.eval()
20
  self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
21
- self.device = "cuda" if torch.cuda.is_available() else "CPU"
22
  self.model.load_state_dict(torch.load(checkpoint_path,map_location=self.device))
23
  self.num_labels = num_labels
24
  self.label_encoder = label_encoder
18
  self.model = model
19
  self.model.eval()
20
  self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
21
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
22
  self.model.load_state_dict(torch.load(checkpoint_path,map_location=self.device))
23
  self.num_labels = num_labels
24
  self.label_encoder = label_encoder