Try to subclass PretrainedModel

modeling_bert.py

@@ -22,7 +22,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-from transformers import PretrainedModel
+from transformers.modeling_utils import PretrainedModel
 from .configuration_bert import JinaBertConfig
 from transformers.models.bert.modeling_bert import (
     BaseModelOutputWithPoolingAndCrossAttentions,
@@ -39,7 +39,6 @@ from flash_attn.modules.block import Block
 from flash_attn.modules.embedding import BertEmbeddings
 from flash_attn.modules.mha import MHA
 from flash_attn.modules.mlp import FusedMLP, Mlp
-from flash_attn.utils.pretrained import state_dict_from_pretrained
 
 try:
     from flash_attn.ops.fused_dense import FusedDense