Overview

PyTorch 2.6, released on January 30, 2025, improves torch.compile and training performance.

Main Features

Improved torch.compile

torch.compile supports more operations and reduces initial compilation time. The reduce-overhead mode is more stable.

python 
import torch

model = torch.nn.Sequential(
    torch.nn.Linear(784, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10),
)

# Compile for acceleration
compiled = torch.compile(model, mode='reduce-overhead')
x = torch.randn(32, 784)
out = compiled(x)
print(f'Output shape: {out.shape}')

Performance improvements

CUDA operations are optimized with better FlashAttention support and fused kernels for Transformers.

python 
import torch
import torch.nn.functional as F

# Scaled dot-product attention (FlashAttention)
q = torch.randn(2, 8, 128, 64)
k = torch.randn(2, 8, 128, 64)
v = torch.randn(2, 8, 128, 64)

out = F.scaled_dot_product_attention(q, k, v)
print(f'Attention shape: {out.shape}')

Sources