Spaces:
Sleeping
Sleeping
set weights only and map location for torch.load
Browse files
app.py
CHANGED
|
@@ -12,7 +12,7 @@ class_names = ["pizza", "steak", "sushi"]
|
|
| 12 |
effnetb0, effnetb0_transforms = create_effnetb0_model()
|
| 13 |
|
| 14 |
effnetb0.load_state_dict(torch.load(
|
| 15 |
-
"07_effnetb0_data_20_percent_10_epochs.pth"))
|
| 16 |
|
| 17 |
|
| 18 |
def predict(img) -> Tuple[Dict, float]:
|
|
|
|
| 12 |
effnetb0, effnetb0_transforms = create_effnetb0_model()
|
| 13 |
|
| 14 |
effnetb0.load_state_dict(torch.load(
|
| 15 |
+
"07_effnetb0_data_20_percent_10_epochs.pth", weights_only=True, map_location=torch.device('cpu')))
|
| 16 |
|
| 17 |
|
| 18 |
def predict(img) -> Tuple[Dict, float]:
|