Update models/BertTunning.py

models/BertTunning.py

@@ -4,6 +4,7 @@ from models.preprocess_stage.bert_model import model
 from models.preprocess_stage.bert_model import preprocess_bert
 
 MAX_LEN = 100
+DEVICE='cpu'
 
 class BertTunnig(nn.Module):
     def __init__(self, bert_model):
@@ -40,6 +41,7 @@ def predict_2(text):
     preprocessed_text, attention_mask = preprocess_bert(text, MAX_LEN=MAX_LEN)
     preprocessed_text, attention_mask = torch.tensor(preprocessed_text).unsqueeze(0), torch.tensor([attention_mask])
 
+    model_tunning.to(DEVICE)
     with torch.inference_mode():
 
         predict = round(model_tunning(preprocessed_text, attention_mask=attention_mask).item())