Update main_test_SRMNet.py
Browse files- main_test_SRMNet.py +1 -1
main_test_SRMNet.py
CHANGED
|
@@ -71,7 +71,7 @@ def save_img(filepath, img):
|
|
| 71 |
|
| 72 |
|
| 73 |
def load_checkpoint(model, weights):
|
| 74 |
-
checkpoint = torch.load(weights)
|
| 75 |
try:
|
| 76 |
model.load_state_dict(checkpoint["state_dict"])
|
| 77 |
except:
|
|
|
|
| 71 |
|
| 72 |
|
| 73 |
def load_checkpoint(model, weights):
|
| 74 |
+
checkpoint = torch.load(weights, map_location=torch.device('cpu'))
|
| 75 |
try:
|
| 76 |
model.load_state_dict(checkpoint["state_dict"])
|
| 77 |
except:
|