🐛 Bugfix: TensorRT output types (#395)

lyuwenyu · Jul 29, 2024 · 6b36128 · 6b36128
1 parent b8957b3
commit 6b36128
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rtdetrv2_pytorch/references/deploy/rtdetrv2_tensorrt.py b/rtdetrv2_pytorch/references/deploy/rtdetrv2_tensorrt.py
@@ -225,6 +225,6 @@ def draw(images, labels, boxes, scores, thrh = 0.6):
         'orig_target_sizes': orig_size.to(args.device),
     }
 
-    labels, boxes, scores = m(blob)
+    output = m(blob)
 
-    draw([im_pil], labels, boxes, scores)
+    draw([im_pil], output['labels'], output['boxes'], output['scores'])