Update app.py

app.py

@@ -191,9 +191,14 @@ class ViT(nn.Module):
 
 # Load model weights
 checkpoint = torch.load("best_ViT_CIFAR100_baseline_checkpoint.pth", map_location=device)
+
+model = ViT(cfg).to(device)
+
+# Load only the model weights
 model.load_state_dict(checkpoint["model_state"])
 model.eval()
 
+
 # Image preprocessing
 transform = transforms.Compose([
     transforms.Resize((32,32)),