Merge pull request #93 from ngshya/patch-1

added map_location=str(device) in torch.load()
sgrvinod · Feb 13, 2020 · 6e61fc6 · 6e61fc6
2 parents 6b0f830 + 094f0b9
commit 6e61fc6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/caption.py b/caption.py
@@ -197,7 +197,7 @@ def visualize_att(image_path, seq, alphas, rev_word_map, smooth=True):
     args = parser.parse_args()
 
     # Load model
-    checkpoint = torch.load(args.model)
+    checkpoint = torch.load(args.model, map_location=str(device))
     decoder = checkpoint['decoder']
     decoder = decoder.to(device)
     decoder.eval()