fix issue when using cpu mode for testing

Author: backseason

joint_main.py

@@ -61,7 +61,7 @@ def main(config):
     parser.add_argument('--n_color', type=int, default=3)
     parser.add_argument('--lr', type=float, default=5e-5) # Learning rate resnet:5e-5, vgg:1e-4
     parser.add_argument('--wd', type=float, default=0.0005) # Weight decay
-    parser.add_argument('--cuda', type=bool, default=True)
+    parser.add_argument('--no-cuda', dest='cuda', action='store_false')
 
     # Training settings
     parser.add_argument('--arch', type=str, default='resnet') # resnet or vgg

joint_solver.py

@@ -25,7 +25,10 @@ def __init__(self, train_loader, test_loader, config):
         self.build_model()
         if config.mode == 'test':
             print('Loading pre-trained model from %s...' % self.config.model)
-            self.net.load_state_dict(torch.load(self.config.model))
+            if self.config.cuda:
+                self.net.load_state_dict(torch.load(self.config.model))
+            else:
+                self.net.load_state_dict(torch.load(self.config.model, map_location='cpu'))
             self.net.eval()
 
     # print the network information and parameter numbers

main.py

@@ -59,7 +59,7 @@ def main(config):
     parser.add_argument('--n_color', type=int, default=3)
     parser.add_argument('--lr', type=float, default=5e-5) # Learning rate resnet:5e-5, vgg:1e-4
     parser.add_argument('--wd', type=float, default=0.0005) # Weight decay
-    parser.add_argument('--cuda', type=bool, default=True)
+    parser.add_argument('--no-cuda', dest='cuda', action='store_false')
 
     # Training settings
     parser.add_argument('--arch', type=str, default='resnet') # resnet or vgg

solver.py

@@ -25,7 +25,10 @@ def __init__(self, train_loader, test_loader, config):
         self.build_model()
         if config.mode == 'test':
             print('Loading pre-trained model from %s...' % self.config.model)
-            self.net.load_state_dict(torch.load(self.config.model))
+            if self.config.cuda:
+                self.net.load_state_dict(torch.load(self.config.model))
+            else:
+                self.net.load_state_dict(torch.load(self.config.model, map_location='cpu'))
             self.net.eval()
 
     # print the network information and parameter numbers