Dimensionality of CNN and a linear function

Hi I need help understanding how the nn.linear can be implemented in a neural network without problems dimensionality.

EDIT

I think my problem relates to understanding the relationship between in and output channel.

Why does this work:

m2 = nn.Linear(32, 31)


input = torch.randn(128, 32)
output = m2(input)
print(output.size())

When this doesn't work:

m2 = nn.Linear(512, 4)
input = torch.randn(128, 2, 2, 4)
output = m2(input)
print(output.size())

Considering that the input of Linear matches the three first dimensions of the tensor, which seems to correspond to the description in the documentation: https://pytorch.org/docs/stable/generated/torch.nn.Linear.html