Update README.md

README.md

@@ -79,19 +79,21 @@ The following hyperparameters were used during training:
 
 ### Code to Run
 
-from transformers import ViTFeatureExtractor
-from transformers import ViTForImageClassification
-import torch
-
-vit = ViTForImageClassification.from_pretrained("oschamp/vit-artworkclassifier")
-vit.eval()
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-vit.to(device)
-
-model_name_or_path = 'google/vit-base-patch16-224-in21k'
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path)
-
 def vit_classify(image):
+    from transformers import ViTFeatureExtractor
+    from transformers import ViTForImageClassification
+    import torch
+    
+    vit = ViTForImageClassification.from_pretrained("oschamp/vit-artworkclassifier")
+    vit.eval()
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    vit.to(device)
+    
+    model_name_or_path = 'google/vit-base-patch16-224-in21k'
+    feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path)
+
+    #LOAD IMAGE
+
     encoding = feature_extractor(images=image, return_tensors="pt")
     encoding.keys()