search

THU-MIG / yolov10

YOLOv10: Real-Time End-to-End Object Detection [NeurIPS 2024]
https://arxiv.org/abs/2405.14458
GNU Affero General Public License v3.0
9.58k stars 913 forks source link

fix issue : 'dict' object has no attribute 'shape' when predict in PyTorch model #273

Closed qianyue76 closed 3 months ago

qianyue76 commented 3 months ago

A small fix with non_max_suppression in ops.py . The fix originated from the ultralytics issue 13587 issues:13587

yingliu0518 commented 2 months ago

I get the same error when I do the inference step. Could you please provide more details about the issue you have done, thank you very much!

I add the comited code to the nms methods, but still encounter this error:'tuple' object has no attribute 'shape'

yingliu0518 commented 2 months ago

This issus was solved by a little change:

if isinstance(prediction, dict) and 'one2one' in prediction:  # model.pt, output = {"one2many": torch.Tensor, "one2one": torch.Tensor }
    prediction = prediction['one2one'][0]