1. Follow this instruction from NVIDIA to install CUDA drivers for WSL.

DO NOT install CUDA Toolkit yet. The version in that instruction is not compatible with current version of Tensorflow/Pytorch.

2. Install CUDA Toolkit 11.8.0 for WSL:

wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-wsl-ubuntu-11-8-local_11.8.0-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-11-8-local_11.8.0-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

Add CUDA to path:

echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.zshrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.zshrc

Symlink libcuda (needed for building darknet):

sudo ln -s /usr/lib/wsl/lib/libcuda.so.1 /usr/local/cuda/lib64/libcuda.so

3. Install cuDNN

Go to https://developer.nvidia.com/rdp/cudnn-download and download cuDNN for CUDA 11.x for Linux.
Filename: cudnn-linux-x86_64-8.9.4.25_cuda11-archive.tar.xz

tar -xvf cudnn-linux-x86_64-8.9.4.25_cuda11-archive.tar.xz
cd cudnn-linux-x86_64-8.9.4.25_cuda11-archive
sudo cp -P include/* /usr/local/cuda/include/
sudo cp -P lib/* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
sudo ldconfig /usr/local/cuda/lib64

4. Install pyenv and create a virtual environment

5. Install Tensorflow and Pytorch:

pip install --upgrade pip setuptools wheel
pip install torch torchvision torchaudio
pip install tensorflow

6. Install additional packages (optional):

pip install jupyterlab