Update model

model.py

@@ -1,4 +1,4 @@
-from transformers import AutoModel
+from transformers import AutoModel, AutoConfig, PreTrainedModel
 import torch
 
 
@@ -15,7 +15,7 @@ class MultiLabelAttention(torch.nn.Module):
         return torch.matmul(torch.transpose(attention_weights, 2, 1), x)
 
 
-class BertMesh(torch.nn.Module):
+class BertMesh(PreTrainedModel):
     def __init__(
         self,
         pretrained_model,
@@ -24,7 +24,8 @@ class BertMesh(torch.nn.Module):
         dropout=0,
         multilabel_attention=False,
     ):
-        super().__init__()
+        super().__init__(config=AutoConfig.from_pretrained(pretrained_model))
+        self.config.auto_map = {"AutoModel": "transformers_model.BertMesh"}
         self.pretrained_model = pretrained_model
         self.num_labels = num_labels
         self.hidden_size = hidden_size