diff --git a/han/sent_level_rnn.py b/han/sent_level_rnn.py
index c0349ef..01ad004 100644
--- a/han/sent_level_rnn.py
+++ b/han/sent_level_rnn.py
@@ -23,7 +23,7 @@ class SentLevelRNN(nn.Module):
             sentence_h,_ = self.sentence_GRU(x)
             x = torch.tanh(self.sentence_linear(sentence_h))
             x = torch.matmul(x, self.sentence_context_wghts)
-            x = x.squeeze()
+            x = x.squeeze(dim=2)
             x = self.soft_sent(x.transpose(1,0))
             x = torch.mul(sentence_h.permute(2,0,1), x.transpose(1,0))
             x = torch.sum(x,dim = 1).transpose(1,0).unsqueeze(0)
diff --git a/han/word_level_rnn.py b/han/word_level_rnn.py
index c6aad7d..0f98f0a 100644
--- a/han/word_level_rnn.py
+++ b/han/word_level_rnn.py
@@ -41,7 +41,7 @@ class WordLevelRNN(nn.Module):
         h,_ = self.GRU(x)
         x = torch.tanh(self.linear(h))
         x = torch.matmul(x, self.word_context_wghts)
-        x = x.squeeze()
+        x = x.squeeze(dim=2)
         x = self.soft_word(x.transpose(1,0))
         x = torch.mul(h.permute(2,0,1), x.transpose(1,0))
         x = torch.sum(x, dim = 1).transpose(1,0).unsqueeze(0)