How does STE realize in quantification aware training?

How does STE realize in quantification aware training?

You could pass the tensor that you want to quantize through a custom torch.autograd.Function
Here is an example where this is done for the DoReFa quantization: PACT/module.py at 025370908c347ed7b461ccf1db6860a81250f72c · KwangHoonAn/PACT · GitHub

In the forward pass, implement whichever quantization scheme you are interested in. Note that the backward pass is made explicit and simply returns the upstream gradient as the gradient for the operation’s input tensor.

Thanks for your reply.