diff --git a/inference.py b/inference.py
index e12392e..df74802 100644
--- a/inference.py
+++ b/inference.py
@@ -117,7 +117,7 @@ def rescale_bboxes(out_bbox, size):
                 'to', 'under', 'using', 'walking in', 'walking on', 'watching', 'wearing', 'wears', 'with']
 
     model, _, _ = build_model(args)
-    ckpt = torch.load(args.resume)
+    ckpt = torch.load(args.resume, map_location=torch.device(args.device))
     model.load_state_dict(ckpt['model'])
     model.eval()