Overview

PyTorch 2.5, released on October 17, 2024, adds CuDNN support and regional compilation to partially optimize models.

Main Features

CuDNN and regional compilation

The CuDNN backend is better integrated and regional compilation allows compiling only specific parts of the model with torch.compile.

python 
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = torch.nn.Linear(256, 128)
        self.decoder = torch.nn.Linear(128, 256)

    def forward(self, x):
        # Regional compilation: only encoder is compiled
        x = torch.compile(self.encoder)(x)
        return self.decoder(x)

model = MyModel()
print(model(torch.randn(1, 256)).shape)

Sources