Ahsen Khaliq commited on
Commit
8f1570d
·
1 Parent(s): 2711d4d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -15,7 +15,7 @@ device = "cpu"
15
  # Pick a pretrained model
16
  model_name = "omnivore_swinB"
17
  torch.hub._validate_not_a_forked_repo=lambda a,b,c: True
18
- model = torch.hub.load("facebookresearch/omnivore", model=model_name)
19
 
20
  # Set to eval mode and move to desired device
21
  model = model.to(device)
 
15
  # Pick a pretrained model
16
  model_name = "omnivore_swinB"
17
  torch.hub._validate_not_a_forked_repo=lambda a,b,c: True
18
+ model = torch.hub.load("facebookresearch/omnivore:main", model=model_name,pretrained=True)
19
 
20
  # Set to eval mode and move to desired device
21
  model = model.to(device)