Update vit_model_test.py

vit_model_test.py

@@ -29,13 +29,14 @@ if __name__ == "__main__":
     # Check for GPU availability
     device = torch.device('cuda')
 
+    #this code runs only with nvidia gpu  
     # Load the pre-trained ViT model and move it to GPU
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(device)
 
 
 
     model.classifier = nn.Linear(model.config.hidden_size, 2).to(device)
-    # Define the image preprocessing pipeline
+    # resize image and make it a tensor (add dimension)
     preprocess = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor()