Update app.py

app.py

@@ -13,8 +13,10 @@ from transformers import ViTForImageClassification, ViTImageProcessor
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load ViT model from Hugging Face
-model_path = "andromeda01111/ViT_BCC/data.pkl"
-vit_model = torch.load(model_path, map_location=torch.device("cpu"))
+model_path = "andromeda01111/ViT_BCC"  # Ensure this file is uploaded to the Space
+state_dict = torch.load(model_path, map_location=device)
+vit_model.load_state_dict(state_dict)
+vit_model.to(device)
 vit_model.eval()
 
 # Load Neural Network model from Hugging Face