@ptrblck I removed batch_first=True and here are all the shapes
    def forward(self, x):
        embeddings = self.word_embedding(x)
        embeddings = self.dropout(embeddings)
        # print(f"Embeddings Shape: {embeddings.shape}")  # torch.Size([50, 128, 300])
        output, (hidden_state, _) = self.lstm(embeddings)
        # print(f"Output Shape: {embeddings.shape}")  # torch.Size([50, 128, 300])
        # print(f"Hidden_state Shape: {hidden_state.shape}")  torch.Size([4, 128, 256])
        encoder_out = output[:, :, :self.hidden_dim] + output[:, :, self.hidden_dim:]
        # print(f"encoder_out Shape: {encoder_out.shape}")  # torch.Size([50, 128, 256])
        encoder_out = encoder_out.permute(1, 0, 2)
        # print(f"encoder_out Shape after premute: {encoder_out.shape}")  # torch.Size([128, 50, 256])
        encoder_hidden = (hidden_state[-2, :, :] + hidden_state[-1, :, :]).unsqueeze(0)
        # print(f"encoder hidden Shape: {encoder_hidden.shape}")  # torch.Size([1, 128, 256])
        attention_output = self.attnetwork(encoder_out, encoder_hidden)
        # print(f"Attention Output Shape: {attention_output.shape}")  # torch.Size([128, 256])
        logits = self.classifier(attention_output)
        # print(f"Logits Shape: {logits.shape}")  # torch.Size([128, 35])
        return logits

I guess the output of the attention layer should be 3D torch.Size([128, 50, 35]) as for every token of the 50 I have to get probability distribution of which label it should belong it, correct me if I am wrong please