The swapaxes operation will create an input of [N, C, W, H], while [N, C, H, W] is expected.
I would recommend to use x = x.permute(0, 3, 1, 2) on the tensor.
This shouldn’t make a huge difference if you are flattening the tensor anyway.



 erm:

Is it possible that using the default Kaiming initialization for nn.Linear changes the results significantly?


Yes, that could make a huge difference.
Could you apply the same initialization as is used in the Coursera code?
I also assume that the model architecture is equal to the one from the course?
The single output feature of hidden looks a bit “odd”, but might still work.