Update app.py

app.py

@@ -13,8 +13,9 @@ from transformers import ViTForImageClassification, ViTImageProcessor
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load ViT model from Hugging Face
-vit_model = ViTForImageClassification.from_pretrained("andromeda01111/ViT_BCC").to(device)
-vit_processor = ViTImageProcessor.from_pretrained("andromeda01111/ViT_BCC")
+model_path = "andromeda01111/ViT_BCC/model.pkl"
+vit_model = torch.load(model_path, map_location=torch.device("cpu"))
+vit_model.eval()
 
 # Load Neural Network model from Hugging Face
 nn_model = tf.keras.models.load_model("andromeda01111/NN_BC")