from transformers import PretrainedConfig

class Resnet50Config(PretrainedConfig):
    # since we have an image classification task
    # we need to put a model type that is close to our task
    # don't worry this will not affect our model
    #model_type = "MobileNetV1"
    def __init__(
        self,
        num_classes=6,
        **kwargs):

        self.num_classes = num_classes
        super().__init__(**kwargs)