Hi Blade!



 blade:

Unfortunately, my numpy reimplementation blows up after few iteration while the pytorch implementation starts to converge.


I would start by checking, step by step, whether your pytorch and
numpy implementations give the same results (up to round-off error).
First initialize your pytorch and numpy models with the same weights.
Then pass the same single sample through both models and see if
the model outputs are the same.  Then check that the loss-function
values are the same.  Then backpropagate both and see if they
generate the same gradients.  Finally take one optimizer step and
see whether you get the same updated weights for both models.
Presumably the two implementations differ somewhere, and this
should show you where.
As an aside, you have four linear layers without any intervening
non-linear “activations.”  As such, the four layers collapse into a
single linear layer.  This won’t break anything, but the example you
are studying is just a complicated way of writing a single linear
layer, Linear (784, 1, bias = False).  Recognizing this might
make it simpler to carry out your debugging.
Best.
K. Frank