Update app.py

app.py

@@ -91,11 +91,11 @@ class PretrainedBertClass(torch.nn.Module):
     def __init__(self):
         super(PretrainedBertClass, self).__init__()
         self.l1 = BertModel.from_pretrained(bert_path)
-        self.l2 = torch.nn.Dropout(0.3)
+        self.l2 = torch.nn.Dropout(HEAD_DROP_OUT)
         self.l3 = torch.nn.Linear(768, 6)
     
     def forward(self, ids, mask, token_type_ids):
-        _, output_1= self.l1(ids, attention_mask = mask, token_type_ids = token_type_ids)
+        _, output_1= self.l1(ids, attention_mask = mask, token_type_ids = token_type_ids, return_dict=False)
         output_2 = self.l2(output_1)
         output = self.l3(output_2)
         return output