From 3dd3ced6616b6a19ae84ec925e3a300b3a0963ef Mon Sep 17 00:00:00 2001 From: Ashutosh-Adhikari Date: Fri, 9 Nov 2018 21:37:49 -0500 Subject: [PATCH] Fix HAN for batch_size 1 (#161) --- han/sent_level_rnn.py | 2 +- han/word_level_rnn.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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)