Setting Up NVIDIA Drivers and CUDA for ML/DL on Ubuntu 22.04

Abraham Audu — Mon, 06 Apr 2026 20:57:28 +0000

Let's face it, Windows is getting really stressful to work with as a developer in 2026. If you've tried to work with your GPU on TensorFlow for deep learning projects, you probably have discovered that you're either stuck with CPU-based compute, dated CUDA versions for older TensorFlow versions, or the hell that is setting up WSL2.

Why not just switch to a Linux native environment like Ubuntu? Yeah, that's easily the safe and sane choice in 2026. Moreover, most of your production workloads will be on Linux servers, so might as well make the jump on your local setup too.

I've struggled in time past to configure all the moving parts. This is because the information is scattered across and and people really do be saying random things on the internet. It can easily take hours to figure it out.

And before you run off to ChatGeePeeDee to show you all the steps, I need you to understand that LLMs hallucinate and Linux will allow you run ANY command, including "removing the French language pack". So be careful what you run on your system from LLMs.

This is an attempt to make it easy, especially as more people dump Windows for Linux for the sake of sanity away from the MicroSlop ecosystem.

This setup was done for an Nvidia 3060 Laptop edition GPU (Use this as context to safety-check the steps as you go along).

Let's just jump right in!

Nvidia Drivers (v595) and CUDA 12.1 Setup for Ubuntu 22.04 x86

Flush old installation

sudo nvidia-uninstall
sudo apt purge -y '^nvidia-*' '^libnvidia-*'
sudo rm -r /var/lib/dkms/nvidia
sudo apt -y autoremove
sudo update-initramfs -c -k `uname -r`
sudo update-grub2
read -p "Press any key to reboot... " -n1 -s
sudo reboot

Update Packages

sudo apt update

Install Kernel Headers and Build Tools

sudo apt install -y linux-headers-$(uname -r) build-essential dkms

Install Nvidia Cuda Keyring

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update

Pin Driver Branch (595 in this case, update as necessary)

sudo apt install -y nvidia-driver-pinning-595

Install the Driver

sudo apt install -y nvidia-open

Reboot System

sudo reboot

Verify Installation

nvidia-smi

You now have NVIDIA Drivers setup, Let's proceed to CUDA and cuDNN installation

TL;DR: CUDA 12.1 with cuDNN 8.9.x is the most stable version for TensorFlow and PyTorch simultaneously, so it's generally best to go with this version combination for now.

Flush Old CUDA Installation

sudo apt remove --purge -y 'cuda*' 'libcudnn*'
sudo rm -rf /usr/local/cuda*

Install CUDA with NVIDIA Network Repo

# Download CUDA 12.1 package keyring
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update

# Install CUDA 12.1 toolkit
sudo apt install -y cuda-toolkit-12-1

Install cuDNN

Download cuDNN 8.9 for CUDA 12.1 from NVIDIA Developer:

[https://developer.nvidia.com/rdp/cudnn-archive]
Choose Linux x86_64 / tar package for CUDA 12.1
Extract and copy files

tar -xvf download-path/cudnn-file-name.tar.xz
cd extracted-folder-path
sudo cp -P include/cudnn*.h /usr/local/cuda/include/
sudo cp -P lib/libcudnn* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

Set Environment variables

nano ~/.bashrc

Replace or Paste these environment variables for CUDA:

export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

Apply Changes

source ~/.bashrc

Reboot

sudo reboot

Verfiy installations

nvcc --version
echo $CUDA_HOME
echo $PATH
echo $LD_LIBRARY_PATH

N.B: nvidia-smi will show the highest compatible CUDA version based on the installed NVIDIA driver version, whilst nvcc --version will show the currently installed CUDA version, so if they differ, this is totally fine.

Verify visibility to python frameworks

Create a python virtual environment
Paste this in a .py file:

import tensorflow as tf
print(tf.config.list_physical_devices('GPU'))

Output: [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

Conclusion

If you made it this far and followed all the steps, you should now be able to run your TensorFlow and PyTorch (and indeed other GPU-based workloads) on your NVIDIA GPU in your Ubuntu environment.

Follow for more tech content around machine learning and data science.

DEV Community: Abraham Audu