Spaces:
Runtime error
Runtime error
Commit
•
5523f82
1
Parent(s):
809a7b1
Update app.py
Browse files
app.py
CHANGED
|
@@ -34,7 +34,7 @@ def load_model(model_type):
|
|
| 34 |
model = U2NET(3,1)
|
| 35 |
if model_type == "U2Net":
|
| 36 |
model_path = "weights/u2net.pth"
|
| 37 |
-
model.load_state_dict(torch.load(model_path))
|
| 38 |
|
| 39 |
else:
|
| 40 |
model_path = "weights/quant_model_u2net.pth"
|
|
|
|
| 34 |
model = U2NET(3,1)
|
| 35 |
if model_type == "U2Net":
|
| 36 |
model_path = "weights/u2net.pth"
|
| 37 |
+
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
|
| 38 |
|
| 39 |
else:
|
| 40 |
model_path = "weights/quant_model_u2net.pth"
|