Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -11,8 +11,12 @@ class_names = ["pizza", "steak", "sushi"]
|
|
| 11 |
|
| 12 |
effnetb2, effnetb2_transform = create_effnetb2_model(num_classes=len(class_names))
|
| 13 |
|
| 14 |
-
|
| 15 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 16 |
|
| 17 |
def predict(img) -> Tuple[Dict, float]:
|
| 18 |
start_time = timer()
|
|
|
|
| 11 |
|
| 12 |
effnetb2, effnetb2_transform = create_effnetb2_model(num_classes=len(class_names))
|
| 13 |
|
| 14 |
+
state_dict = torch.load(
|
| 15 |
+
"pretrained_effnetb2_pizza_steak_sushi_20_percent.pt",
|
| 16 |
+
map_location=torch.device("cpu") # Load model to the CPU
|
| 17 |
+
)
|
| 18 |
+
|
| 19 |
+
effnetb2.load_state_dict(state_dict)
|
| 20 |
|
| 21 |
def predict(img) -> Tuple[Dict, float]:
|
| 22 |
start_time = timer()
|