Update vit_model_test.py

vit_model_test.py

@@ -7,13 +7,15 @@ from PIL import Image
 
 class CustomModel:
     def __init__(self):
-        # Check for GPU availability
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # Explicitly set the device to CPU
+        self.device = torch.device('cpu')
 
-        # Load the pre-trained ViT model and move it to GPU
+        # Load the pre-trained ViT model and move it to CPU
         self.model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(self.device)
         self.model.classifier = nn.Linear(self.model.config.hidden_size, 2).to(self.device)
-        self.model.load_state_dict(torch.load('trained_model.pth'))
+
+        # Load model weights
+        self.model.load_state_dict(torch.load('trained_model.pth', map_location=self.device))
         self.model.eval()
 
         # Define the image preprocessing pipeline