Merge pull request #1476 from johnnynunez/patch-1

Update generate_mot_results.py
mikel-brostrom · Jun 12, 2024 · 4a0cd73 · 4a0cd73
2 parents d16a9a3 + 6e8ba8f
commit 4a0cd73
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tracking/generate_mot_results.py b/tracking/generate_mot_results.py
@@ -22,18 +22,19 @@
 from ultralytics.data.utils import VID_FORMATS
 
 from tracking.utils import convert_to_mot_format, write_mot_results
-
+from boxmot.utils.torch_utils import select_device
 
 __tr = TestRequirements()
 __tr.check_packages(('ultralytics @ git+https://github.com/mikel-brostrom/ultralytics.git', ))  # install
 
 
 def generate_mot_results(args):
+    args.device = select_device(args.device)
     tracker = create_tracker(
         args.tracking_method,
         TRACKER_CONFIGS / (args.tracking_method + '.yaml'),
         args.reid_model.with_suffix('.pt'),
-        'cpu',
+        args.device,
         False,
         False
     )