In the code class WeightedPermuteMLP, after h = h.permute(0, 3, 2, 1, 4), h becomes (B, segment_dim, W, H, S); after w = w.permute(0, 1, 3, 2, 4), w becomes (B, H, segment_dim, W, S). My question is why can't w also be turned into (B, segment_dim, H, W, S)

[lxy51]

In the code class WeightedPermuteMLP, after h = h.permute(0, 3, 2, 1, 4), h becomes (B, segment_dim, W, H, S); after w = w.permute(0, 1, 3, 2, 4), w becomes (B, H, segment_dim, W, S). My question is why can't w also be turned into (B, segment_dim, H, W, S)