使用yolov5训练时报错该怎么解决？

Traceback (most recent call last):
  File "train.py", line 669, in <module>
    main(opt)
  File "train.py", line 564, in main
    train(opt.hyp, opt, device, callbacks)
  File "train.py", line 400, in train
    compute_loss=compute_loss)
  File "/opt/conda/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
    return func(*args, **kwargs)
  File "/root/yolov5-master/val.py", line 266, in run
    stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy   
  File "/root/yolov5-master/val.py", line 266, in <listcomp>
    stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy   
RuntimeError: All input tensors must be on the same device. Received cpu and cuda:0

前面几轮训练都很正常的