Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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):
|