Update app.py

app.py

@@ -19,7 +19,11 @@ os.system("wget https://huggingface.co/akhaliq/Omnivore/resolve/main/swinB_check
 # Pick a pretrained model 
 model_name = "omnivore_swinB"
 model = torch.hub.load('facebookresearch/omnivore:main', "omnivore_swinB", pretrained=False)
-model.load_state_dict(torch.load('/home/user/app/swinB_checkpoint.torch')['trunk'])
+new_dict = {}
+for key, value in torch.load('/home/user/app/swinB_checkpoint.torch')['trunk'].items():
+    new_dict['trunk.' + key] = value
+
+model.load_state_dict(new_dict)
 
 # Set to eval mode and move to desired device
 model = model.to(device)