Update augvit_config.py

augvit_config.py

@@ -16,6 +16,7 @@ class AugViTConfig(PretrainedConfig):
         mlp_dim: int = 256,
         dropout: int = 0.1,
         emb_dropout: int = 0.1,
+        num_channels:int=3,
         **kwargs,
     ):
 
@@ -28,4 +29,5 @@ class AugViTConfig(PretrainedConfig):
         self.mlp_dim = mlp_dim
         self.dropout = dropout
         self.emb_dropout = emb_dropout
+        self.num_channels=num_channels
         super().__init__(**kwargs)