Update dnabert_layer.py

dnabert_layer.py

@@ -5,6 +5,7 @@ import torch.nn as nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.models.bert.modeling_bert import BertModel as TransformersBertModel
 from transformers.models.bert.modeling_bert import BertForMaskedLM as TransformersBertForMaskedLM
+from transformers.models.bert.modeling_bert import BertForPreTraining as TransformersBertForPreTraining
 from transformers.models.bert.modeling_bert import BertPreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput
 
@@ -17,6 +18,11 @@ class BertForMaskedLM(TransformersBertForMaskedLM):
     def __init__(self, config):
         super().__init__(config)
 
+class BertForPreTraining(TransformersBertForPreTraining):
+    def __init__(self, config):
+        super().__init__(config)
+
+
 
 class DNABertForSequenceClassification(BertPreTrainedModel):
     def __init__(self, config):