Update app.py

app.py

@@ -19,7 +19,7 @@ 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'))
+model.load_state_dict(torch.load('/home/user/app/swinB_checkpoint.torch')['state_dict'])
 
 # Set to eval mode and move to desired device
 model = model.to(device)