jojojo:

What was the motivation to choose this interpretation and choosing dim 4 as sequence dimension?


One idea could be that the spatial dimensions are comparable to a temporal dimension in that both are encoding a sequential information (i.e. neighboring values might be correlated).



 jojojo:

Shouldn’t it be W = T?


Yes, you are right. My comments are wrong, T = x.size(1) should still work as it would assign W to T.



 jojojo:

Also, shouldn’t input to linear layer be in shape of 2D tensor? Shouldn’t we flatten it? You have passed a 3D tensor to linear layer. Will it work?


Yes, it will work as described in the docs. Linear layers accept an input of [batch_size, *, in_features] where * denotes additional dimensions. They will be applied on each sample in the additional dimensions as if you were using a for loop.