my_config.py

from transformers import PretrainedConfig


class MyModelConfig(PretrainedConfig):
    model_type = "my_model"

    def __init__(
        self,
        n_layers=2,
        **kwargs,
    ):
        self.n_layers = n_layers
        super().__init__(**kwargs)