modeling.py

from transformers import BertModel

from .configuration import NewModelConfig
from .util import ImportLayer

class NewModel(BertModel):
    config_class = NewModelConfig

    def __init__(self, config):
        super().__init__(config)
        self.last_layer = ImportLayer(config.hidden_size, config.new_hidden_size)