CatoEr commited on
Commit
50239d2
1 Parent(s): 7216ad1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -41,11 +41,11 @@ orientation_labels = {
41
 
42
  model_race = RaceClassifier(n_classes=4)
43
  model_race.to(device)
44
- model_race.load_state_dict(torch.load('best_model_race.pt'))
45
 
46
  model_orientation = RaceClassifier(n_classes=2)
47
  model_orientation.to(device)
48
- model_orientation.load_state_dict(torch.load('best_model_orientation_last.pt'))
49
 
50
 
51
  def evaluate(model, input, mask):
 
41
 
42
  model_race = RaceClassifier(n_classes=4)
43
  model_race.to(device)
44
+ model_race.load_state_dict(torch.load('best_model_race.pt', map_location=torch.device('cpu')))
45
 
46
  model_orientation = RaceClassifier(n_classes=2)
47
  model_orientation.to(device)
48
+ model_orientation.load_state_dict(torch.load('best_model_orientation_last.pt', map_location=torch.device('cpu')))
49
 
50
 
51
  def evaluate(model, input, mask):