Enable PyTorch with DirectML on WSL

PyTorch with DirectML enables training and inference on DirectX 12-capable GPUs in Windows Subsystem for Linux (WSL). PyTorch with DirectML is in public preview and works in WSL 2.

Check your version of Windows

The torch-directml package in WSL 2 requires Windows 11, build 22000 or later. To check your Windows version and build number, select Windows logo key + R, enter winver, and select OK.

Install WSL 2

To install the default Linux distribution with WSL 2, open PowerShell or Windows Command Prompt in administrator mode and run:

wsl --install

Restart your machine when prompted. For distribution selection and other installation options, see Install Linux on Windows with WSL.

Check for GPU driver updates

Install the latest Windows driver available for your GPU through Windows Update or your hardware manufacturer's website. The Windows driver enables GPU acceleration in WSL; you don't need to install a separate Linux display driver.

Set up Python

Install a Python environment in your WSL distribution. For example, run the following commands to install Miniconda:

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

Then create and activate an environment named pytorch-directml:

conda create --name pytorch-directml python=3.10
conda activate pytorch-directml

Install PyTorch with DirectML

Install the torch-directml package:

pip install torch-directml

Verify the installation

Start Python and run the following code to add two tensors on the DirectML device:

import torch
import torch_directml

dml = torch_directml.device()
tensor1 = torch.tensor([1]).to(dml)
tensor2 = torch.tensor([2]).to(dml)
dml_algebra = tensor1 + tensor2
print(dml_algebra.item())

Expected output:

3

Samples and feedback

See the DirectML PyTorch samples for examples. To report package issues or request features, use the DirectML issue tracker.