彻底解决NotImplementedError: Could not run ‘torchvision::nms‘ with arguments from the ‘CUDA‘ backend.

原因：

安装的是torchvision torchaudio的cpu版本，用GPU训练某些模型时会存在问题，例如下面的报错：

NotImplementedError: Could not run 'torchvision::roi_align' with arguments from the 'CUDA' backend.

NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend.

解决方案：

因为pytorch官网提供的安装命令安装的是cpu版本的torchvision和torchaudio，如果想安装cuda版本的，可以使用下面的命令：

pip install torch==2.3.1 torchaudio==2.3.1 torchvision==0.18.1 -f https://download.pytorch.org/whl/cu118/torch_stable.html

把上面的版本号替换为自己想要的版本，注意-f参数之后的网址中的cu118替换为自己的cuda版本。

安装完成后，pip list，应该是类似这样的：

torch 2.3.1+cu118

torchaudio 2.3.1+cu118

torchvision 0.18.1+cu118

这样就不会出现torchvision::nms 和 torchvision::roi_align等问题了，一劳永逸。