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

change model weight save path

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -40,8 +40,9 @@ model.fc = nn.Sequential(
40
  nn.Linear(model.fc.in_features, 2)
41
  )
42
 
 
43
  urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
44
- model.load_state_dict(torch.load(MODEL_PATH))
45
  model.eval()
46
  for _, p in model.named_parameters():
47
  p.requires_grad = False
 
40
  nn.Linear(model.fc.in_features, 2)
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