Update app.py

app.py

@@ -35,7 +35,7 @@ class CustomViTModel(nn.Module):
 
 # Load the trained model
 model = CustomViTModel()
-model.load_state_dict(torch.load('final_ms_model_2classes_vit_base_224_bce.pth', map_location=device))
+model.load_state_dict(torch.load('final_ms_model_2classes_vit_base_384_bce.pth', map_location=device))
 model.to(device)
 model.eval()