Added as custom model

config.json

@@ -1,6 +1,6 @@
 {
   "architectures": [
-    "Bert2DForTokenClassification"
+    "Bert2DForMaskedLM"
   ],
   "attention_probs_dropout_prob": 0.1,
   "gradient_checkpointing": false,
@@ -20,5 +20,10 @@
   "position_embedding_type": "absolute",
   "type_vocab_size": 2,
   "use_cache": false,
-  "vocab_size": 128000
+  "vocab_size": 128000,
+  "auto_map": {
+    "AutoConfig": "configuration_bert2d.Bert2DConfig",
+    "AutoModelForMaskedLM": "modeling_bert2d.Bert2DForMaskedLM",
+    "AutoTokenizer": "tokenization_bert2d_fast.Bert2DTokenizerFast"
+  }
 }