From 557bfaddd3062905609673b3323e7230e7316ecd Mon Sep 17 00:00:00 2001 From: carloscaetano Date: Wed, 26 Feb 2025 14:19:49 -0300 Subject: [PATCH] adding CPU device when loading the model --- inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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()