Enable PyTorch with DirectML on Windows

PyTorch with DirectML enables training and inference on DirectX 12-capable GPUs. PyTorch with DirectML is in public preview and works on native Windows starting with Windows 10 version 1709.

Check your version of Windows

To check your Windows version and build number, select Windows logo key + R, enter winver, and select OK. Update to Windows 10 version 1709 or later if your build is older.

Check for GPU driver updates

Install the latest driver available for your GPU through Windows Update or your hardware manufacturer's website.

Set up Python

Install a Python environment. If you use Miniconda, download and run the Windows installer for your architecture.

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.