Update app.py

app.py

@@ -190,8 +190,8 @@ class ViT(nn.Module):
 
 
 # Load model weights
-model = ViT(cfg).to(device)
-model.load_state_dict(torch.load("best_ViT_CIFAR100_baseline_checkpoint.pth", map_location=device))
+checkpoint = torch.load("best_ViT_CIFAR100_baseline_checkpoint.pth", map_location=device)
+model.load_state_dict(checkpoint["model_state"])
 model.eval()
 
 # Image preprocessing