from transformers import PretrainedConfig

class SqueezeNetConfig(PretrainedConfig):
    model_type = "squeezenet_custom"

    def __init__(self, num_classes: int = 200, fire_dropout: float = 0.0, final_dropout: float = 0.50, **kwargs):
        self.num_classes = num_classes
        self.fire_dropout = fire_dropout
        self.final_dropout = final_dropout
        super().__init__(**kwargs)