Assume there is a non-differentiable nn.Parameter in an equation whose gradient needs to be estimated using a straight-through estimator (STE) before the parameter can be updated. For example,
y = 0.5(|x| - |x - alpha| + alpha)
y_q = round(y * (2 ** k - 1) / alpha) * alpha / (2 ** k - 1).
In this equation, alpha is a trainable parameter and the derivative of y_q w.r.t alpha needs to be estimated using an STE.
How can I define the approximated gradient and use it in an optimizer to update the parameter along with the rest of parameters in the network?