Update cuda and cudnn for pop os

 

Update the nvidia-driver

sudo apt purge '^nvidia-*'
sudo apt autoremove
sudo apt clean
sudo apt update
sudo apt full-upgrade
sudo apt install system76-driver-nvidia # check the version of driver
sudo systemctl reboot

Install the cuda and cudnn 12.6

Now, the cuda and cudnn can be installed via pip.

  1. Create a python env with conda
conda create --name cu126 python=3.10
  1. To install Wheels, you must first install the nvidia-pyindex package, which is required in order to set up your pip installation to fetch additional Python modules from the NVIDIA NGC PyPI repo. If your pip and setuptools Python modules are not up-to-date, then use the following command to upgrade these Python modules. If these Python modules are out-of-date then the commands which follow later in this section may fail.
python3 -m pip install --upgrade setuptools pip wheel
  1. You should now be able to install the nvidia-pyindex module.
python3 -m pip install nvidia-pyindex
  1. Install the following metapackages (may unnecessary for installing pytorch)
python3 -m pip install nvidia-cuda-runtime-cu12
python3 -m pip install nvidia-cuda-cupti-cu12
python3 -m pip install nvidia-cuda-nvcc-cu12
python3 -m pip install nvidia-nvml-dev-cu12
python3 -m pip install nvidia-cuda-nvrtc-cu12
python3 -m pip install nvidia-nvtx-cu12
python3 -m pip install nvidia-cuda-sanitizer-api-cu12
python3 -m pip install nvidia-cublas-cu12
python3 -m pip install nvidia-cufft-cu12
python3 -m pip install nvidia-curand-cu12
python3 -m pip install nvidia-cusolver-cu12
python3 -m pip install nvidia-cusparse-cu12
python3 -m pip install nvidia-npp-cu12
python3 -m pip install nvidia-nvjpeg-cu12
python3 -m pip install nvidia-opencl-cu12
python3 -m pip install nvidia-nvjitlink-cu12
python3 -m pip install nvidia-cudnn-cu12
  1. Install pytorch, e.g. 2.6
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

One can replace the https://download.pytorch.org/whl with https://mirrors.aliyun.com/pytorch-wheels/, as mirrored by aliyun

pip install torch torchvision torchaudio --index-url https://mirrors.aliyun.com/pytorch-wheels/cu126
  1. Validate the cuda version pytorch is using
python
>>>torch
>>>print(torch.version.cuda)