Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -17,7 +17,7 @@ model = UNet(
|
|
17 |
num_res_units=4,
|
18 |
dropout=0.15,
|
19 |
)
|
20 |
-
model.load_state_dict(torch.load("best_model.pth"
|
21 |
model.eval()
|
22 |
|
23 |
def greet(image_path):
|
|
|
17 |
num_res_units=4,
|
18 |
dropout=0.15,
|
19 |
)
|
20 |
+
model.load_state_dict(torch.load("best_model.pth", map_location=torch.device('cpu')))
|
21 |
model.eval()
|
22 |
|
23 |
def greet(image_path):
|