Hi, why do we need to swap axis before we flatten?

[diff7]

Hi, I can't figure out why we need to change from BCHW to BHWC before we flatten. I would be happy if you could explain this moment. Thank you!

convert inputs from BCHW -> BHWC

    inputs = inputs.permute(0, 2, 3, 1).contiguous()
    input_shape = inputs.shape
    
    # Flatten input
    flat_input = inputs.view(-1, self._embedding_dim)