The problem is that you are just directly reshaping the tensors to the shape you want. That does not work, because it “scrambles” the data. You need to do it in two steps:
- Reshape from (n_H, n_W, n_C) to (n_H * n_W, n_C).
- Then transpose that to get the final result.
They give hints about this in the instructions and then in the pictures, but they aren’t as explicit about it as perhaps they should have been. The reason for this is similar to why the flattening of the images in Course 1 Week 2 requires the transpose. Here’s a thread about that particular case. The high level point is that just because the shapes end up the same does not mean the results are equivalent. It matters exactly how you do it.
Update: here’s a more recent thread that shows the analogous experiment to the previous link, but it’s specific to the case here in C4 W4 A2. The link points to the experiment, but the earlier parts of thread are worth a look as well.