PyTorch 2.1: Dynamic Shapes and torch.export

Published on 05/10/2023 86 views pytorch ml release

PyTorch 2.1 stabilizes dynamic shapes in torch.compile and introduces torch.export.

Overview

PyTorch 2.1, released on October 5, 2023, stabilizes dynamic shapes in torch.compile and introduces torch.export for model export.

Main Features

Dynamic shapes let torch.compile handle variable-size tensors without recompilation. torch.export produces a portable IR graph.

python

import torch

@torch.compile(dynamic=True)
def f(x):
    return x.sin() + x.cos()

print(f(torch.randn(4)))
print(f(torch.randn(8)))  # no recompilation

Sources

PyTorch 2.1 Release