I’m not sure, how all variables are defined, e.g. self.b, but you could try to remove certain parts of the code, which involve the linear layer for its calculation, and try to isolate which operation might decrease the gradient magnitude.
If I understand the code correctly, you are applying a manual sigmoid, which might create small gradients, if the input is saturated (very small or very large).