Update app.py

app.py

@@ -22,6 +22,9 @@ model = torch.hub.load('facebookresearch/omnivore:main', "omnivore_swinB", pretr
 new_dict = {}
 for key, value in torch.load('/home/user/app/swinB_checkpoint.torch')['trunk'].items():
     new_dict['trunk.' + key] = value
+  
+for key, value in torch.load('/home/user/app/swinB_checkpoint.torch')['heads'].items():
+    new_dict['heads.' + key] = value
 
 model.load_state_dict(new_dict)