config.py

from transformers import PretrainedConfig
class MNIST_config(PretrainedConfig):
    model_type = "MNIST_Classifier"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
  
        for key,value in kwargs.items():
            setattr(self,key,value)
        #print(key,value)
        #self.input_size=kwargs['input_size']
        #self.hidden_size1=kwargs["hidden_size1"]
        #self.hidden_size2=kwargs["hidden_size2"]
        #self.output_size=kwargs["output_size"]