Hi,

You should apply the same logic to .bias if you want it to work as well.
You should NOT modify the original parameters inplace. That is why you have the weight_orig that are the original weights that are optimized. and the weight that are the ones that are the processed version of weight_orig after your kmeans like op and that is used in the forward.
Keep in mind that you’re doing differentiable operations. And that if you override a value, the gradient for the original value before the override is 0.