MNCJihun commited on
Commit
67e9c3a
·
1 Parent(s): 582ad59

change model weight save path

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -13,6 +13,7 @@ import urllib.request
13
 
14
  MODEL_URL = "https://huggingface.co/caisarl76/HI_motorcycle_trunk_cls_model/resolve/main/best_model.pth"
15
  MODEL_PATH = "/tmp/best_model.pth"
 
16
 
17
  test_transforms = A.Compose(
18
  [
@@ -41,8 +42,13 @@ model.fc = nn.Sequential(
41
  )
42
 
43
  device = torch.device('cpu')
44
- urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
45
- model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
 
 
 
 
 
46
  model.eval()
47
  for _, p in model.named_parameters():
48
  p.requires_grad = False
 
13
 
14
  MODEL_URL = "https://huggingface.co/caisarl76/HI_motorcycle_trunk_cls_model/resolve/main/best_model.pth"
15
  MODEL_PATH = "/tmp/best_model.pth"
16
+ urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
17
 
18
  test_transforms = A.Compose(
19
  [
 
42
  )
43
 
44
  device = torch.device('cpu')
45
+
46
+ try:
47
+ model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
48
+ print('model load complete')
49
+ except:
50
+ print('CANNOT load model weight')
51
+
52
  model.eval()
53
  for _, p in model.named_parameters():
54
  p.requires_grad = False