Charig Yang
commited on
Commit
·
5195bf1
1
Parent(s):
3f0fdc0
up
Browse files
app.py
CHANGED
|
@@ -11,7 +11,7 @@ def predict(img):
|
|
| 11 |
model = models.resnet50()
|
| 12 |
model.fc = nn.Linear(2048, 720)
|
| 13 |
resume_path = 'full+++++.pth'
|
| 14 |
-
model.load_state_dict(torch.load(resume_path))
|
| 15 |
model.to(device)
|
| 16 |
with torch.no_grad():
|
| 17 |
model.eval()
|
|
|
|
| 11 |
model = models.resnet50()
|
| 12 |
model.fc = nn.Linear(2048, 720)
|
| 13 |
resume_path = 'full+++++.pth'
|
| 14 |
+
model.load_state_dict(torch.load(resume_path, map_location=torch.device(device)))
|
| 15 |
model.to(device)
|
| 16 |
with torch.no_grad():
|
| 17 |
model.eval()
|