Installing WSL2, PyTorch and CUDA on Windows 11

If you have a compatible Nvidia GPU, you can work seamlessly with Ubuntu Linux and CUDA within your regular Windows 11 (22H2) OS. Really, it’s the best of both worlds! However installing the frameworks and tools is really not that well documented, so hopefully this walkthrough will help.

The first step is to make sure that you are running the latest version of the Nvidia driver, available from https://www.nvidia.com/Download/index.aspx?lang=en-us. Note, you no longer need a specific driver from Nvidia, as you have previously.

Next we need to install WSL2 and Ubuntu. This step is very straight-forward, just fire up a PowerShell and enter:

wsl --install

When entered, Windows will prompt you for elevated privileges, and then proceed to install the WSL frameworks and Ubuntu. At this point, I ran into the following well-documented error:

Ubuntu is already installed.
Launching Ubuntu...
Installing, this may take a few minutes...
WslRegisterDistribution failed with error: 0x80370114
Error: 0x80370114 The operation could not be started because a required feature is not installed.

I solved this by clicking on the search icon, and opening “Turn Windows Features On and Off” and switching on Windows Subsystem For Linux. Following a restart, I then ran:

 wsl.exe --install ubuntu

Ubuntu will then launch, and prompt you to set a username and password, and you will then be greeted by the bash prompt. Congratulations, you now have WSL2 installed!

Now a neat thing, is if you pull up the Windows Terminal app, you can hit the drop-down arrow on the right of the tabs, and launch directly into an Ubuntu instance. Pretty cool!

Now within WSL2, as is pretty standard with any new Ubuntu installation, I ran the following to update the system:

sudo apt update && sudo apt upgrade

Now I like to work with minimal bloat, so contrary to other reports online, I went ahead and skipped Anaconda and started by making sure that I had both Python and Pip installed:

sudo apt install python3 python3-pip

I then went to the PyTorch website and installed the latest version (not nightly) for Linux, CUDA 11.8 and Pip. The command it generated for me was:

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

This process installs Torch and all of the necessary libraries. It took a few minutes to complete, but when complete you should now have a working PyTorch installation within WSL2 with CUDA support enabled!

Now to validate that everything went smoothly, enter a Python environment with python3 and enter the following commands:

>>> import torch
>>> torch.cuda.is_available()
True
>>>

You should see True as shown above. If so, you can now proceed to build upon this minimal installation by installing whichever packages you need, for example Jupyter, Numpy and Matplotlib.

Finally, I am also aware that Microsoft is working quite extensively on their DirectML APIs, and now support CUDA in WSL2 over DirectML. My experience has been varied with this, so have chosen to take the direct CUDA route. Keen to hear everyones comments!