From be128d4d9f7b187bbd2697bb44151359f06c48b3 Mon Sep 17 00:00:00 2001 From: akarazniewicz Date: Mon, 11 Nov 2019 12:36:06 +0100 Subject: [PATCH] Mapping loaded, pretrained model to actual device. Required for CPU only inference. --- demo.py | 2 +- test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/demo.py b/demo.py index ed679dd..71d059d 100755 --- 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 index a215825..c18abe7 100755 --- 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)