Update README.md

README.md

@@ -161,11 +161,12 @@ class CustomViTForImageClassification(ViTForImageClassification):
             attentions=outputs.attentions,
         )
 
-
-model = CustomViTForImageClassification.from_pretrained("vesteinn/vit-mae-cub")
-image_processor = AutoImageProcessor.from_pretrained("vesteinn/vit-mae-cub")
-
-classifier = pipeline(
-    "image-classification", model=model, image_processor=image_processor
-)
+if __name__ == "__main__":
+  
+  model = CustomViTForImageClassification.from_pretrained("vesteinn/vit-mae-inat21")
+  image_processor = AutoImageProcessor.from_pretrained("vesteinn/vit-mae-inat21")
+  
+  classifier = pipeline(
+      "image-classification", model=model, image_processor=image_processor
+  )
 ```