Update modeling_me2bert.py

modeling_me2bert.py

@@ -1,9 +1,8 @@
-from transformers import PretrainedConfig
 from transformers import PreTrainedModel
 from transformers import AutoModel
 import torch
 from torch.autograd import Function
-
+from .configuration_me2bert import ME2BertConfig
 
 class ReverseLayerF(Function):
 
@@ -30,8 +29,8 @@ class FFClassifier(torch.nn.Module):
             torch.nn.BatchNorm1d(hidden_dim), torch.nn.ReLU(True),
             torch.nn.Dropout(dropout), torch.nn.Linear(hidden_dim, n_classes))
 
-    def forward(self, input):
-        return self.model(input)
+    def forward(self, x):
+        return self.model(x)
 
 
 class Encoder(torch.nn.Module):
@@ -102,7 +101,7 @@ class GatedCombination(torch.nn.Module):
 
 class ME2BertModel(PreTrainedModel):
     config_class = ME2BertConfig
-
+    base_model_prefix = "me2bert"
     def __init__(
             self,
             config: ME2BertConfig = None):