Error about detect.py

When I execute the command to test "detect.py",

python detect.py
--source test.mp4
--weights yolov9-e.pt
--imgsz 640
--conf-thres 0.25
--iou-thres 0.1

I encountered the following error:

Traceback (most recent call last):
File "/home/gvsai/yolov9/detect.py", line 231, in
main(opt)
File "/home/gvsai/yolov9/detect.py", line 226, in main
run(**vars(opt))
File "/home/gvsai/anaconda3/envs/yolov5/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/home/gvsai/yolov9/detect.py", line 102, in run
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/gvsai/yolov9/utils/general.py", line 905, in non_max_suppression
device = prediction.device
^^^^^^^^^^^^^^^^^
AttributeError: 'list' object has no attribute 'device'

The following line needs to be changed:

yolov9/detect.py

Line 98 in 9660d12

pred = model(im, augment=augment, visualize=visualize)

to:

pred, _ = model(im, augment=augment, visualize=visualize)

Then "detect.py" can be executed normally!

#11

@mkrupczak3
Thanks!!