Fix bug in dimensions

Author: emalgorithm

models/lstm_decoder.py

@@ -18,13 +18,14 @@ def __init__(self, hidden_size, output_size, device, attention=False):
         self.device = device
 
     def forward(self, input, hidden, encoder_hiddens):
+        # encoder_hiddens has shape [batch_size, seq_len, hidden_dim]
         output = self.embedding(input).view(1, 1, -1)
 
         if self.attention:
             hiddens = torch.cat((encoder_hiddens, hidden[0].repeat(1, encoder_hiddens.size(1), 1)),
-                                dim=1)
+                                dim=2)
             attention_coeff = self.attention_layer(hiddens)
-            context = torch.mm(torch.squeeze(encoder_hiddens, dim=1).t(), torch.squeeze(
+            context = torch.mm(torch.squeeze(encoder_hiddens, dim=0).t(), torch.squeeze(
                 attention_coeff, 2).t()).view(1, 1, -1)
             output = torch.cat((output, context), 2)
             output = self.attention_combine(output)