pt_custom_model_example / my_config.py
ydshieh
add config
75b4b63
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)