Hugging Face's logo

Spaces:
Amould
/
2DPcloudreg
Runtime error

App Files Community
Amould commited on
Commit
6224444
·
verified ·
1 Parent(s): e03c91e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -29,10 +29,10 @@ file_savingfolder = './modelbest/'
29
  ext = '_bestVal'
30
  core_model_tst = resnet18(pretrained=True)
31
  core_model_tst.fc = Identity()
32
- core_model_tst.load_state_dict(torch.load(file_savingfolder+'core_model'+ext+'.pth'))
33
  #core_model_tst.to(device)
34
  IR_Model_tst = Build_IRmodel_Resnet(core_model_tst, registration_method)
35
- IR_Model_tst.load_state_dict(torch.load(file_savingfolder+'IR_Model'+ext+'.pth'))
36
  #IR_Model_tst.to(device)
37
  IR_Model_tst.eval()
38
 
 
29
  ext = '_bestVal'
30
  core_model_tst = resnet18(pretrained=True)
31
  core_model_tst.fc = Identity()
32
+ core_model_tst.load_state_dict(torch.load(file_savingfolder+'core_model'+ext+'.pth', map_location=torch.device('cpu') ))
33
  #core_model_tst.to(device)
34
  IR_Model_tst = Build_IRmodel_Resnet(core_model_tst, registration_method)
35
+ IR_Model_tst.load_state_dict(torch.load(file_savingfolder+'IR_Model'+ext+'.pth', map_location=torch.device('cpu')))
36
  #IR_Model_tst.to(device)
37
  IR_Model_tst.eval()
38