Update vit_model_test.py

vit_model_test.py

@@ -15,7 +15,7 @@ class CustomModel:
         self.model.classifier = nn.Linear(self.model.config.hidden_size, 2).to(self.device)
 
         # Load model weights
-        self.model.load_state_dict(torch.load('trained_model.pth', map_location=self.device))
+        self.model.load_state_dict(torch.load('trained_model.pth', map_location=self.device, weights_only=True))
         self.model.eval()
 
         # Define the image preprocessing pipeline