Bug fix

Author: LLXXTTT

eval.py

@@ -106,7 +106,8 @@ def eval(args):
     torch.set_grad_enabled(False)
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = get_model(args.model, args.dataset)
-    model_state = torch.load(args.checkpoint)['model_state']
+    model_state = torch.load(args.checkpoint, 
+                map_location=device)['model_state']
     model.load_state_dict(model_state)
     model.to(device)
     model.eval()

train.py

@@ -62,7 +62,7 @@ def train(args):
         if os.path.isfile(args.resume):
             print("Loading model and optimizer from checkpoint '{}'".format\
                   (args.resume))
-            checkpoint = torch.load(args.resume)
+            checkpoint = torch.load(args.resume, map_location=device)
             model.load_state_dict(checkpoint['model_state'])
             optimizer.load_state_dict(checkpoint['optimizer_state'])
             print("Loaded checkpoint '{}' (epoch{})".format(args.resume,