No, I do not think clip_grad will work, what I want to do is, zero out elements of a matrix, that fall below a threshold, but this threshold is not fixed, it is learnable, I do not set it, model needs to learn this threshold, and zero out elements of matrix that fall below this threshold.