Update predict.py

Add the process checking whether GPU or CPU can be used

predict.py

@@ -19,7 +19,7 @@ from bs4 import BeautifulSoup
 # import pandas as pd
 # import numpy as np
 # import codecs
-
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
 #%% global変数として使う
 dict_key = {}
@@ -60,7 +60,7 @@ def predict_entities(modelpath, sentences_list, len_num_entity_type):
     # bert_tc = model.bert_tc.cuda()
     
     model = ner.BertForTokenClassification_pl(modelpath, num_labels=81, lr=1e-5) 
-    bert_tc = model.bert_tc.cuda()
+    bert_tc = model.bert_tc.to(device)
 
     MODEL_NAME = 'cl-tohoku/bert-base-japanese-whole-word-masking'
     tokenizer = ner.NER_tokenizer_BIO.from_pretrained(
@@ -79,7 +79,7 @@ def predict_entities(modelpath, sentences_list, len_num_entity_type):
             encoding, spans = tokenizer.encode_plus_untagged(
                 text, return_tensors='pt'
             )
-            encoding = { k: v.cuda() for k, v in encoding.items() } 
+            encoding = { k: v.to(device) for k, v in encoding.items() } 
             
             with torch.no_grad():
                 output = bert_tc(**encoding)