iscyy / yoloair

🔥🔥🔥 专注于YOLOv5,YOLOv7、YOLOv8、YOLOv9改进模型,Support to improve backbone, neck, head, loss, IoU, NMS and other modules🚀
https://github.com/iscyy/yoloair
GNU General Public License v3.0
2.47k stars 428 forks source link

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) #159

Open AKA-WDNMD opened 1 year ago

AKA-WDNMD commented 1 year ago

Describe the bug RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

To Reproduce Steps to reproduce the behavior: 1.运行python train.py --img 640 --batch 4 --epochs 50 --data coco.yaml --weights yolov5s.pt --device '0'

ecfa0b573078b0cda62bcc5b8c4a085

解决方法:

  1. ./utils/loss.py的685行改为 :from_which_layer.append((torch.ones(size=(len(b),)) * i).to('cuda'))
  2. ./utils/loss.py的756行之后添加多一行:fg_mask_inboxes = fg_mask_inboxes.to(torch.device('cuda'))
    参考:https://stackoverflow.com/questions/74372636/indices-should-be-either-on-cpu-or-on-the-same-device-as-the-indexed-tensor