Skip to content

Commit 73f7cc3

Browse files
oToToTcopybara-github
authored andcommitted
Update PyTorch PT2E tutorial for PyTorch 2.6+
PyTorch 2.6+ no lnoger has `torch._export.capture_pre_autograd_graph`. We update the sample code here and the pytorch tutorial link to reflect the change PiperOrigin-RevId: 842522823
1 parent 5f08e65 commit 73f7cc3

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

docs/pytorch_converter/README.md

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -120,13 +120,11 @@ edge_model(*inputs_2, signature_name="input2")
120120
## Quantization
121121

122122
Following is the code snippet to quantize a model with [PT2E
123-
quantization](https://pytorch.org/tutorials/prototype/quantization_in_pytorch_2_0_export_tutorial.html)
123+
quantization](https://docs.pytorch.org/ao/stable/tutorials_source/pt2e_quant_ptq.html)
124124
using the `ai_edge_torch` backend.
125125

126126
```python
127127
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
128-
from torch._export import capture_pre_autograd_graph
129-
130128
from ai_edge_torch.quantize.pt2e_quantizer import get_symmetric_quantization_config
131129
from ai_edge_torch.quantize.pt2e_quantizer import PT2EQuantizer
132130
from ai_edge_torch.quantize.quant_config import QuantConfig
@@ -135,7 +133,12 @@ pt2e_quantizer = PT2EQuantizer().set_global(
135133
get_symmetric_quantization_config(is_per_channel=True, is_dynamic=True)
136134
)
137135

138-
pt2e_torch_model = capture_pre_autograd_graph(torch_model, sample_args)
136+
# > For pytorch 2.6+
137+
pt2e_torch_model = torch.export.export(torch_model, sample_args).module()
138+
# > For pytorch 2.5 and before
139+
# from torch._export import capture_pre_autograd_graph
140+
# pt2e_torch_model = capture_pre_autograd_graph(torch_model, sample_args)
141+
139142
pt2e_torch_model = prepare_pt2e(pt2e_torch_model, pt2e_quantizer)
140143

141144
# Run the prepared model with sample input data to ensure that internal observers are populated with correct values

0 commit comments

Comments
 (0)