yoloair
yoloair copied to clipboard
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
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'
解决方法:
- ./utils/loss.py的685行改为 :from_which_layer.append((torch.ones(size=(len(b),)) * i).to('cuda'))
- ./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