Update app.py
Browse files
app.py
CHANGED
|
@@ -35,7 +35,7 @@ labels = {
|
|
| 35 |
}
|
| 36 |
model_race = RaceClassifier(n_classes=4)
|
| 37 |
model_race.to(device)
|
| 38 |
-
model_race.load_state_dict(torch.load('best_model_race.pt'
|
| 39 |
|
| 40 |
|
| 41 |
def predict(text):
|
|
|
|
| 35 |
}
|
| 36 |
model_race = RaceClassifier(n_classes=4)
|
| 37 |
model_race.to(device)
|
| 38 |
+
model_race.load_state_dict(torch.load('best_model_race.pt', map_location=torch.device('cpu')))
|
| 39 |
|
| 40 |
|
| 41 |
def predict(text):
|