I am trying to add attention mechanism to stacked LSTMs implementation https://github.com/salesforce/awd-lstm-lm 232
All examples online use encoder-decoder architecture, which I do not want to use (do I have to for the attention mechanism?).
Basically, I have used https://webcache.googleusercontent.com/search?q=cache:81Q7u36DRPIJ:https://github.com/zhedongzheng/finch/blob/master/nlp-models/pytorch/rnn_attn_text_clf.py+&cd=2&hl=en&ct=clnk&gl=uk 191
    class RNNModel(nn.Module):
        def __init__(self, rnn_type, ntoken, ninp, nhid, nlayers, dropout=0.5, dropouth=0.5, dropouti=0.5, dropoute=0.1, wdrop=0, tie_weights=False):
            super(RNNModel, self).__init__()
            self.encoder = nn.Embedding(ntoken, ninp)
                self.rnns = [torch.nn.LSTM(ninp if l == 0 else nhid, nhid if l != nlayers - 1 else (ninp if tie_weights else nhid), 1, dropout=0) for l in range(nlayers)]
                for rnn in self.rnns:
                    rnn.linear = WeightDrop(rnn.linear, ['weight'], dropout=wdrop)
            self.rnns = torch.nn.ModuleList(self.rnns)
            self.attn_fc = torch.nn.Linear(ninp, 1)
            self.decoder = nn.Linear(nhid, ntoken)
    
            self.init_weights()
    
        def attention(self, rnn_out, state):
            state = state.squeeze(0).unsqueeze(2)
            weights = torch.bmm(rnn_out, state)
            weights = torch.nn.functional.softmax(weights.squeeze(2)).unsqueeze(2)
            # (batch, cell_size, seq_len) * (batch, seq_len, 1) = (batch, cell_size, 1)
            return torch.bmm(torch.transpose(rnn_out, 1, 2), weights).squeeze(2)
    
        def forward(self, input, hidden, return_h=False, decoder=False, encoder_outputs=None):
            emb = embedded_dropout(self.encoder, input, dropout=self.dropoute if self.training else 0)
            emb = self.lockdrop(emb, self.dropouti)
    
            raw_output = emb
            new_hidden = []
            raw_outputs = []
            outputs = []
            for l, rnn in enumerate(self.rnns):
                current_input = raw_output
                raw_output, new_h = rnn(raw_output, hidden[l])
                new_hidden.append(new_h)
                raw_outputs.append(raw_output)
                if l != self.nlayers - 1:
                    raw_output = self.lockdrop(raw_output, self.dropouth)
                    outputs.append(raw_output)
            hidden = new_hidden
    
            output = self.lockdrop(raw_output, self.dropout)
            outputs.append(output)
    
            attn_out = self.attention(raw_output, new_h[0])
    
            decoded = self.decoder(attn_out)
            return decoded

input has 5 items of length 5
model.forward returns torch.Size([5, 831])
Unfortunately, this implementation of attention is missing one dimension (i.e. 5x5x831)
How can I modify attention function to bring back the dimension?