Update bert_layers.py

bert_layers.py

@@ -18,7 +18,8 @@ from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
 from transformers.activations import ACT2FN
 from transformers.modeling_outputs import (MaskedLMOutput,
                                            SequenceClassifierOutput)
-from transformers.models.bert.modeling_bert import BertPreTrainedModel, PreTrainedModel
+from transformers.models.bert.modeling_bert import BertPreTrainedModel
+from transformers.modeling_utils import PreTrainedModel
 
 from .bert_padding import (index_first_axis,
                                             index_put_first_axis, pad_input,