Mapping loaded, pretrained model to actual device. Required for CPU o…

…nly inference.
WJinhai · Nov 11, 2019 · be128d4 · be128d4
1 parent 603d463
commit be128d4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/demo.py b/demo.py
@@ -30,7 +30,7 @@ def demo(opt):
 
     # load model
     print('loading pretrained model from %s' % opt.saved_model)
-    model.load_state_dict(torch.load(opt.saved_model))
+    model.load_state_dict(torch.load(opt.saved_model, map_location=device))
 
     # prepare data. two demo images from https://github.com/bgshih/crnn#run-demo
     AlignCollate_demo = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)

diff --git a/test.py b/test.py
@@ -165,7 +165,7 @@ def test(opt):
 
     # load model
     print('loading pretrained model from %s' % opt.saved_model)
-    model.load_state_dict(torch.load(opt.saved_model))
+    model.load_state_dict(torch.load(opt.saved_model, map_location=device))
     opt.experiment_name = '_'.join(opt.saved_model.split('/')[1:])
     # print(model)