Bug fix

AaltoVision · Jun 26, 2020 · 543ff9f · 543ff9f
1 parent 2d0c3a5
commit 543ff9f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/eval.py b/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()

diff --git a/train.py b/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,