Nigus commited on
Commit
a793711
1 Parent(s): 0d74fb6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -7
app.py CHANGED
@@ -6,16 +6,13 @@ import os
6
 
7
  torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
8
 
9
- model = torch.hub.load('pytorch/vision:v0.10.0', 'shufflenet_v2_x1_0', pretrained=True)
10
- model.load_state_dict(torch.load('shufflenetv2_x1-5666bf0f80.pth'))
11
-
12
  # To export
13
- model = torch.hub.load('pytorch/vision:v0.10.0', 'shufflenet_v2_x1_0', pretrained=True).eval()
14
- traced_graph = torch.jit.trace(model, torch.randn(1, 3, H, W))
15
- traced_graph.save('shufflenet.pth')
16
 
17
  # To load
18
- model = torch.jit.load('shufflenet.pth').eval().to(device)
19
 
20
 
21
  os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
 
6
 
7
  torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
8
 
 
 
 
9
  # To export
10
+ model = torch.hub.load('pytorch/vision:v0.10.0', 'shufflenet_v2_x1_0', pretrained=True).eval()
11
+ traced_graph = torch.jit.trace(model, torch.randn(1, 3, H, W))
12
+ traced_graph.save('shufflenet.pth')
13
 
14
  # To load
15
+ model = torch.jit.load('shufflenet.pth').eval().to(device)
16
 
17
 
18
  os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")