Yolo v5 官方代码本地跑遇到的问题

出现以下错误提示

Traceback (most recent call last):
  File "train.py", line 470, in <module>
    train(hyp, opt, device, tb_writer)
  File "train.py", line 333, in train
    plots=epoch == 0 or final_epoch)  # plot first and last
  File "/project/train/src_repo/yolov5-master_jishi/test.py", line 193, in test
    plot_images(img, output_to_target(output, width, height), paths, str(f), names)  # predictions
  File "/project/train/src_repo/yolov5-master_jishi/utils/general.py", line 951, in output_to_target
    return np.array(targets)
  File "/usr/local/lib/python3.7/dist-packages/torch/_tensor.py", line 678, in __array__
    return self.numpy()
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

我的解决办法是将utils/general.py文件中

def output_to_target(output, width, height):
    # Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
    if isinstance(output, torch.Tensor):
        output = output.cpu().numpy()

    targets = []
    for i, o in enumerate(output):
        if o is not None:
            for pred in o:
                box = pred[:4]
                w = (box[2] - box[0]) / width
                h = (box[3] - box[1]) / height
                x = box[0] / width + w / 2
                y = box[1] / height + h / 2
                conf = pred[4]
                cls = int(pred[5])
                # 列表里边有tensor数据，需要.cpu()
                targets.append([i, cls, x.cpu(), y.cpu(), w.cpu(), h.cpu(), conf.cpu()])

    return np.array(targets)

targets.append([i, cls, x.cpu(), y.cpu(), w.cpu(), h.cpu(), conf.cpu()])
x,y,w,h,conf 这几个tensor全部加上.cpu()，就可以跑通代码了。