add missing class

automodel.py

@@ -6,7 +6,8 @@ import torch.nn as nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
 from transformers import BertPreTrainedModel
-from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers.modeling_outputs import (MaskedLMOutput,
+                                           SequenceClassifierOutput)
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file