Update app.py

app.py

@@ -15,10 +15,11 @@ import gradio as gr
 # Device on which to run the model
 # Set to cuda to load on GPU
 device = "cpu"
-
+os.system("wget https://huggingface.co/akhaliq/Omnivore/resolve/main/swinB_checkpoint.torch")
 # Pick a pretrained model 
 model_name = "omnivore_swinB"
-model = torch.hub.load("facebookresearch/omnivore:main", model=model_name)
+model = torch.hub.load('facebookresearch/omnivore:main', "omnivore_swinB", pretrained=False)
+model.load_state_dict(torch.load('/home/user/app/swinB_checkpoint.torch'))
 
 # Set to eval mode and move to desired device
 model = model.to(device)