# Adding linear layer to transformer model (+ save\_pretrained and load\_pretrained)

**URL:** <https://discuss.huggingface.co/t/adding-linear-layer-to-transformer-model-save-pretrained-and-load-pretrained/15548>\
**Category:** 🤗Transformers\
**Created:** [March 9, 2022, 9:46pm UTC](https://discuss.huggingface.co/t/adding-linear-layer-to-transformer-model-save-pretrained-and-load-pretrained/15548 "2022-03-09T21:46:40Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![waelabid](https://avatars.discourse-cdn.com/v4/letter/w/a9adbd/32.png) [@waelabid](https://discuss.huggingface.co/u/waelabid)\
**Post date:** [March 9, 2022, 9:46pm UTC](https://discuss.huggingface.co/t/adding-linear-layer-to-transformer-model-save-pretrained-and-load-pretrained/15548/1 "2022-03-09T21:46:40Z")

</div>

I want to extend a transformer model (let’s take bert, electra, etc for example) with a linear layer and initialize the linear layer with the same initializer as the transformer model. I also want save\_pretrained and load\_pretrained to work smoothly (i.e. save and load the model WITH the linear layer, not separately). This is what I’m currently doing, but it’s not working

```auto
class ExtendedTransformer(PreTrainedModel):
    base_model_prefix = "model"
    def __init__ (self, config):
        super(). __init__ (config)
        self.model = AutoModel.from_pretrained('bert-base-uncased')
        self.linear = torch.nn.Linear(self.config.hidden_size, 128, bias=False)
        self.model.init_weights()

```

Any advice on how to do this properly?  
Thanks

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [March 10, 2022, 8:37am UTC](https://discuss.huggingface.co/t/adding-linear-layer-to-transformer-model-save-pretrained-and-load-pretrained/15548/2 "2022-03-10T08:37:24Z")

</div>

If you just want to increase the output dimensions, you can simply use

```python
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=128)

```

But here’s an explanation of what I think the issue is with your code.

If I’m not mistaken, your code it only initializes the `model` (which is already being initialized by the `from_pretrained` call). As far as I know, you should sub-class specific architectures, because `PreTrainedModel` itself does not have an `init_weights` method and models may implement a different scheme for that. Since the end of last year, it’s probably best to use `post_init` instead of `init_weights`, too. Maybe you can achieve this for a variety of models with AutoModel as well, but I am not sure.

So, the following should work. It is basically the same as [BertForSequenceClassification](https://github.com/huggingface/transformers/blob/65f9653ed069076fd2b3bdcdfffecbd43c2d2a39/src/transformers/models/bert/modeling_bert.py#L1501-L1515).

```auto
from transformers import BertPreTrainedModel, AutoConfig, AutoModel
import torch

class ExtendedTransformer(BertPreTrainedModel):
    def __init__ (self, config):
        super(). __init__ (config)
        self.bert = AutoModel.from_pretrained('bert-base-uncased')
        self.linear = torch.nn.Linear(self.config.hidden_size, 128, bias=False)
        self.post_init()

if __name__ == ' __main__':
    config = AutoConfig.from_pretrained('bert-base-uncased')
    inst = ExtendedTransformer(config)

```