How to get Tensorflow working with a NVIDIA GPU on Windows using WSL
This seems like it should have been easy based on the instructions from Tensorflow website, but it was not for me. I watched many unhelpful YouTube videos and tried many unhelpful LLM generated suggestions. I finally figured out that the problem was that Tensorflow could not find the NVIDIA libraries it needed. That was brutal, so I wanted to document it clearly for everyone else who might be struggling with the same thing.
This guide assumes
1. You have the NVIDIA GPU that is supported by tensorflow
2. You have correctly installed the NVIDIA driver in Windows
3. You have WSL working in Windows 11.
Here are the steps with commands you can copy and paste:
0. Make sure nvidia driver is working in Windows
nvidia-smi
1. Fresh Ubuntu 24.04 in WSL2
wsl --install Ubuntu
Make sure nvidia driver is working in WSL
nvidai-smi
2. Install python3-venv and python3-pip
sudo apt update
sudo apt install python3-venv python3-pip
3. Create a virtual environment in your Linux home directory (~), not in /mnt/c/
I use tf-gpu as the virtual environment name, you can use whatever you want
cd ~
python3 -m venv tf-gpu
4. Activate the virtual environment
cd tf-gpu
source bin/activate
5. Install tensorflow and CUDA
pip install --upgrade pip
pip install tensorflow[and-cuda]
Verify installation
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
6. Make symbolic links to libraries in correct location
cd $(dirname $(python -c 'print(__import__("tensorflow").__file__)'))
ln -svf ../nvidia/*/lib/*.so* .
cd -
ln -sf $VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cuda_nvcc/bin/ptxas $VIRTUAL_ENV/bin/ptxas
7. Add to LD_LIBRARY_PATH
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$VIRTUAL_ENV/lib/python3.12/site-packages/tensorflow
8. Add LD_LIBRARY_PATH to activate script to make it permanent whenever you use virtual environment