From 708766679d12eecaf0ef69b17c0fb6d43ac5374f Mon Sep 17 00:00:00 2001
From: Amirreza Arjmand <arjmandamirreza@gmail.com>
Date: Wed, 21 Aug 2024 00:28:43 +0330
Subject: [PATCH] minor fix for running inference.py on cpu

---
 inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/inference.py b/inference.py
index e12392e..07b67b8 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('cpu')) if args.device == "cpu" else torch.load(args.resume)
     model.load_state_dict(ckpt['model'])
     model.eval()