Aravindan commited on
Commit
1d1e250
·
1 Parent(s): 7d5de97

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -9,7 +9,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
9
  new_model = models.resnet18()
10
  num_ftrs = new_model.fc.in_features
11
  new_model.fc = nn.Linear(num_ftrs, 75)
12
- checkpoint = torch.load('model_best_checkpoint.pth.tar')
13
  new_model.load_state_dict(checkpoint['model'])
14
  new_model.to(device)
15
 
 
9
  new_model = models.resnet18()
10
  num_ftrs = new_model.fc.in_features
11
  new_model.fc = nn.Linear(num_ftrs, 75)
12
+ checkpoint = torch.load('model_best_checkpoint.pth.tar', map_location=torch.device('cpu'))
13
  new_model.load_state_dict(checkpoint['model'])
14
  new_model.to(device)
15