Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -90,7 +90,7 @@ if __name__ =='__main__':
|
|
90 |
|
91 |
print('Number fo parameters: {}'.format(model.num_parameters()))
|
92 |
model.to(device)
|
93 |
-
model_dict = torch.load('Model_example.pth.tar')['state_dict']
|
94 |
# save model
|
95 |
model.load_state_dict(model_dict, strict=True)
|
96 |
model.eval()
|
|
|
90 |
|
91 |
print('Number fo parameters: {}'.format(model.num_parameters()))
|
92 |
model.to(device)
|
93 |
+
model_dict = torch.load('Model_example.pth.tar', map_location=device)['state_dict']
|
94 |
# save model
|
95 |
model.load_state_dict(model_dict, strict=True)
|
96 |
model.eval()
|