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

[AKA-WDNMD]

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'

解决方法：

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