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.
- Create a python env with conda
conda create --name cu126 python=3.10
- 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
- You should now be able to install the
nvidia-pyindex
module.
python3 -m pip install nvidia-pyindex
- 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
- 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
- Validate the cuda version pytorch is using
python
>>>torch
>>>print(torch.version.cuda)