torchvision::nms 函数无法在 CUDA 后端运行

[weigengming]

我是在算力云平台跑的，但运行完一轮后报错：
Traceback (most recent call last):
File "tools/train.py", line 395, in
main()
File "tools/train.py", line 330, in main
da_segment_results,ll_segment_results,detect_results, total_loss,maps, times = validate(
File "/root/YOLOP-main/lib/core/function.py", line 258, in validate
output = non_max_suppression(inf_out, conf_thres= config.TEST.NMS_CONF_THRESHOLD, iou_thres=config.TEST.NMS_IOU_THRESHOLD, labels=lb)
File "/root/YOLOP-main/lib/core/general.py", line 169, in non_max_suppression
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
File "/root/miniconda3/lib/python3.8/site-packages/torchvision/ops/boxes.py", line 42, in nms
return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
RuntimeError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend. 'torchvision::nms' is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, Tracer, Autocast, Batched, VmapMode].
经搜索后是说函数没法在CUDA后端运行，如何解决

[weigengming]

环境是：
PyTorch 1.7.0
Python 3.8(ubuntu18.04)
Cuda 11.0

[weigengming]

已解决。https://blog.csdn.net/weixin_51179975/article/details/119082342