File tree Expand file tree Collapse file tree 1 file changed +7
-4
lines changed
Expand file tree Collapse file tree 1 file changed +7
-4
lines changed Original file line number Diff line number Diff line change @@ -120,13 +120,11 @@ edge_model(*inputs_2, signature_name="input2")
120120## Quantization
121121
122122Following 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 )
124124using the ` ai_edge_torch ` backend.
125125
126126``` python
127127from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
128- from torch._export import capture_pre_autograd_graph
129-
130128from ai_edge_torch.quantize.pt2e_quantizer import get_symmetric_quantization_config
131129from ai_edge_torch.quantize.pt2e_quantizer import PT2EQuantizer
132130from 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+
139142pt2e_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
You can’t perform that action at this time.
0 commit comments