Spaces:
Running
Running
Hu
commited on
Commit
·
60caba2
1
Parent(s):
7b74848
load with device mapping
Browse files
app.py
CHANGED
@@ -24,7 +24,7 @@ print("Loading SRCNN model...")
|
|
24 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
25 |
|
26 |
model = SRCNNModel().to(device)
|
27 |
-
model.load_state_dict(torch.load('SRCNNmodel_trained.pt'))
|
28 |
model.eval()
|
29 |
print("SRCNN model loaded!")
|
30 |
|
|
|
24 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
25 |
|
26 |
model = SRCNNModel().to(device)
|
27 |
+
model.load_state_dict(torch.load('SRCNNmodel_trained.pt',map_location=torch.device(device) ))
|
28 |
model.eval()
|
29 |
print("SRCNN model loaded!")
|
30 |
|