Update configuration_vit.py

configuration_vit.py

@@ -1 +1,32 @@
- 
+from transformers import PretrainedConfig
+
+class CustomViTNanoV2Config(PretrainedConfig):
+    model_type = "custom_vit_nano"
+
+    def __init__(
+        self,
+        image_size=224,
+        patch_size=16,
+        in_chans=3,
+        num_classes=1000,
+        embed_dim=224,
+        depth=6,
+        num_heads=4,
+        mlp_hidden_dim=608,
+        stem_channels=(32, 64, 128),
+        dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.num_classes = num_classes
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.mlp_hidden_dim = mlp_hidden_dim
+        self.stem_channels = stem_channels
+        self.dropout = dropout
+
+CustomViTNanoV2Config.register_for_auto_class()