It’s also worth having a look at this thread to see the “dimensional analysis” for this test case. It really helps debugging to know what you should be expecting at each layer.
If you look at that, you’ll see that 3 x 4 is the shape of W2 and 5 x 4 is the shape of X. If you are at layer 2, the “linear activation” calculation should be:
Z2 = W2 \cdot A1 + b2
Where the shape of A1 is 4 x 4. So why would you end up with X as the input at layer 2?