small fix

app.py

@@ -19,6 +19,7 @@ class GradioApp:
         custom.load_state_dict(torch.load('models/my_vit.pt', map_location=device))
         
         pretrained = models.vit_b_16().to(device).eval()
+        pretrained.heads = nn.Linear(768, 3)
         pretrained.load_state_dict(torch.load('models/pretrained_vit.pt', map_location=device))
 
         self.models: Dict[str, Union[str, nn.Module]] = {