QuantWrapper — PyTorch 2.7 documentation (original) (raw)

class torch.ao.quantization.QuantWrapper(module)[source][source]¶

A wrapper class that wraps the input module, adds QuantStub and DeQuantStub and surround the call to module with call to quant and dequant modules.

This is used by the quantization utility functions to add the quant and dequant modules, before convert function QuantStub will just be observer, it observes the input tensor, after convert, QuantStubwill be swapped to nnq.Quantize which does actual quantization. Similarly for DeQuantStub.