[bug fixed] load pre-trained checkpoint.

Author: ganji15

eval_demo.py

@@ -31,7 +31,7 @@
     cfg = yaml2config(args.config)
 
     model = get_model(cfg.model)(cfg, args.config)
-    model.load(args.ckpt)
+    model.load(args.ckpt, cfg.device)
     if args.mode == 'style':
         model.eval_style()
     elif args.mode == 'rand':
@@ -41,4 +41,4 @@
     elif args.mode == 'text':
         model.eval_text()
     else:
-        print('Unsupported mode: {} | [rand] [style] [interp]'.format(cfg.mode))
+        print('Unsupported mode: {} | [rand] [style] [interp]'.format(cfg.mode))

networks/model.py

@@ -81,14 +81,18 @@ def save(self, tag='best', epoch_done=0, **kwargs):
         ckpt_save_path = os.path.join(self.log_root, self.opt.training.ckpt_dir, tag + '.pth')
         torch.save(ckpt, ckpt_save_path)
 
-    def load(self, ckpt, modules=None):
+    def load(self, ckpt, map_location=None, modules=None):
         if modules is None:
             modules = []
         elif not isinstance(modules, list):
             modules = [modules]
 
         print('load checkpoint from ', ckpt)
-        ckpt = torch.load(ckpt)
+        if map_location is None:
+            ckpt = torch.load(ckpt)
+        else:
+            ckpt = torch.load(ckpt, map_location=map_location)
+
         if len(modules) == 0:
             for model in self.models.values():
                 model.load_state_dict(ckpt[type(model).__name__])

test.py

@@ -42,5 +42,5 @@
     cfg.valid.dset_split = args.split
 
     model = get_model(cfg.model)(cfg, args.config)
-    model.load(args.ckpt)
-    print(model.validate(args.guided))
+    model.load(args.ckpt, cfg.device)
+    print(model.validate(args.guided))