Update app.py

app.py

@@ -20,7 +20,7 @@ genai.configure(api_key=GEMINI_API_KEY)
 # Load the model
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = VisionTransformer(img_size=128, patch_size=8, num_classes=len(CLASS_NAMES), embed_dim=768, depth=8, num_heads=12, mlp_dim=2048, dropout=0.1)
-model.load_state_dict(torch.load("custom_vit.pth", map_location=DEVICE))
+model.load_state_dict(torch.load("models/custom_vit.pth", map_location=DEVICE))
 model.to(DEVICE)
 model.eval()