Update README.md

README.md

@@ -9,12 +9,5 @@ class CustomViTModel(ViTForImageClassification, PyTorchModelHubMixin):
         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
-)
+model = CustomViTModel.from_pretrained("img2gpspenn/vit")
+model.eval()