nn.Linear expects the input to have the shape [batch_size, *, nb_features], the tensor should not be completely flattened to a 1-dim tensor.
Usually you would use out = out.view(out.size(0), -1) before feeding the activations to the linear layer.
However, I’m currently unsure, if you are using batch_first=True or not, so you might need to permute the dimensions before.