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