How did you ever figure this out?

The weights loaded fine, but as I was training the loss made no sense. So I wrote some basic matmul test.