Update config.json

config.json

@@ -3,7 +3,7 @@
     "ViTForImageClassification"
   ],
   "model": {
-    "type": "timm_wrapper",
+    
     "variant": "vit_small_patch16_384.augreg_in21k_ft_in1k",
     "input_size": 384,
     "patch_size": 16,
@@ -23,6 +23,7 @@
     "resize_size": 440,
     "crop_size": 384
   },
+  "model_type": "vit",
   "dtype": "float32",
   "checkpoint_path": "pretrained_weights/model_v11_ViT_384_base_ckpt.pt"
 }