Update README.md

README.md

@@ -2,3 +2,19 @@
 license: mit
 ---
 lat_mean=39.951163 lat_std=0.000596 lon_mean=-75.19156 lon_std=0.000609
+
+
+class CustomViTModel(ViTForImageClassification, PyTorchModelHubMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        self.classifier = nn.Linear(self.classifier.in_features, config.num_labels)
+
+model_name = "WinKawaks/vit-tiny-patch16-224"
+
+num_classes = 2
+
+custom_model = CustomViTModel.from_pretrained(
+    model_name,
+    num_labels=num_classes,
+    ignore_mismatched_sizes=True
+)