# Loading and save different models types in 1 class

**URL:** <https://discuss.huggingface.co/t/loading-and-save-different-models-types-in-1-class/18862>\
**Category:** 🤗Transformers\
**Created:** [June 8, 2022, 1:50pm UTC](https://discuss.huggingface.co/t/loading-and-save-different-models-types-in-1-class/18862 "2022-06-08T13:50:33Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![shon711](https://avatars.discourse-cdn.com/v4/letter/s/eada6e/32.png) [@shon711](https://discuss.huggingface.co/u/shon711)\
**Post date:** [June 8, 2022, 1:50pm UTC](https://discuss.huggingface.co/t/loading-and-save-different-models-types-in-1-class/18862/1 "2022-06-08T13:50:33Z")

</div>

Hello,

I created my custom class:

```auto
class LingMessCoref(BertPreTrainedModel):
    def __init__ (self, config, args):
        super(). __init__ (config)

        base_model = AutoModel.from_config(config)
        self.base_model_prefix = base_model.base_model_prefix
        self.config_class = base_model.config_class
        setattr(self, self.base_model_prefix, base_model)

        self.clf = Linear(config.hidden_size, 1)

        self.init_weights()

```

I did the `base_model` thing to be able to change the transformer to one of these options: `['longformer', 'roberta', 'deberta', 'bert']`

the first initialization is:

```auto
model, loading_info = LingMessCoref.from_pretrained(
        args.model_name_or_path, output_loading_info=True,
        config=config, cache_dir=args.cache_dir, args=args
    )

```

where `model_name_or_path = 'allenai/longformer-large-4096'`

To save the model checkpoint I run:  
`model.save_pretrained(output_dir)`

Then if I want to load the model again using `from_pretrained` it seems like the weights not loaded.  
the model accuracy and loss is different from the checkpoint run.

Is it make sense that when I change `self` to LingMessCoref it `works` ?

```auto
base_model = AutoModel.from_config(config)
LingMessCoref.base_model_prefix = base_model.base_model_prefix
LingMessCoref.config_class = base_model.config_class
setattr(self, self.base_model_prefix, base_model)

```

Any suggestions?

Thanks in advance,  
Shon