binaries which you can download with
docker pull ghcr.io/pytorch/pytorch-nightly


And for ad hoc experiments just make sure that your container has access
to all your GPUs
docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash


Getting started
a toy exmaple
Let’s start with a simple example and make things more complicated step
by step. Please note that you’re likely to see more significant speedups the newer your GPU is.
import torch
def fn(x, y):
    a = torch.sin(x).cuda()
    b = torch.sin(y).cuda()
    return a + b
new_fn = torch.compile(fn, backend="inductor")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn(input_tensor, input_tensor)

This example won’t actually run faster but it’s educational.