It’s an interesting question that has come up before. It turns out that just getting the correct output shape is not enough: your method actually scrambles the data, because it doesn’t preserve the “samples” dimension. Here’s a thread which explains this in detail.