Return hidden with lengths and masks

uoo723 · Nov 23, 2020 · 11c4d91 · 11c4d91
1 parent 64c13b3
commit 11c4d91
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/deepxml/models.py b/deepxml/models.py
@@ -94,9 +94,9 @@ def train_step(self, train_x: torch.Tensor, train_y: torch.Tensor):
                 emb, train_y = self.mixup_fn(emb, train_y)
                 scores = self.model((emb, lengths, masks), pass_emb=True)
             else:
-                hidden = self.model(train_x, return_hidden=True)
+                hidden, lengths, masks = self.model(train_x, return_hidden=True)
                 hidden, train_y = self.mixup_fn(hidden, train_y)
-                scores = self.model(hidden, pass_hidden=True)
+                scores = self.model((hidden, lengths, masks), pass_hidden=True)
         else:
             scores = self.model(train_x)
 

diff --git a/deepxml/networks.py b/deepxml/networks.py
@@ -65,10 +65,10 @@ def forward(self, inputs, return_emb=False, pass_emb=False,
         if not pass_hidden:
             rnn_out = self.lstm(emb_out, lengths)       # N, L, hidden_size * 2
         else:
-            rnn_out = inputs
+            rnn_out, lengths, masks = inputs
 
         if return_hidden:
-            return rnn_out
+            return rnn_out, lengths, masks
 
         attn_out = self.attention(rnn_out, masks)   # N, labels_num, hidden_size * 2