PyTorch当前支持的CUDA只支持到11.1，你可以通过命令 sudo nvidia-smi查看你当前驱动支持的cuda版本，这里我没法确定版本需不需要强制一定匹配，也没有测试。
我这里按照强制一定匹配的思路来，CUDA 11.1对应的驱动程序是nvidia-driver-455，但是我在ubuntu20上安装这个驱动，总是报失败。
后来看到国外一个大神，安装nvidia-driver-460驱动，其对应的cuda版本是11.2，这个大神说cuda-11.2也可以在pytorch上使用，请参考他的博文