added map_location=str(device) in torch.load()

sgrvinod · Jan 15, 2020 · 094f0b9 · 094f0b9
1 parent b046704
commit 094f0b9
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()