Pack & pad tensors to the LSTM

Author: AngledLuffa

stanza/models/tokenization/model.py

@@ -48,7 +48,9 @@ def forward(self, x, feats, lengths):
         feats = self.dropout_feat(feats)
 
         emb = torch.cat([emb, feats], 2)
+        emb = pack_padded_sequence(emb, lengths, batch_first=True)
         inp, _ = self.rnn(emb)
+        inp, _ = pad_packed_sequence(inp, batch_first=True)
 
         if self.args['conv_res'] is not None:
             conv_input = emb.transpose(1, 2).contiguous()
@@ -75,7 +77,9 @@ def forward(self, x, feats, lengths):
             inp2 = inp
             if self.args['hier_invtemp'] > 0:
                 inp2 = inp2 * (1 - self.toknoise(torch.sigmoid(-tok0 * self.args['hier_invtemp'])))
+            inp2 = pack_padded_sequence(inp2, lengths, batch_first=True)
             inp2, _ = self.rnn2(inp2)
+            inp2, _ = pad_packed_sequence(inp2, batch_first=True)
 
             inp2 = self.dropout(inp2)