b9e27136a5e80850df7859b8f3ca23d396f1b009,snorkel/learning/pytorch/rnn/lstm.py,LSTM,forward,#LSTM#Any#Any#,22

Before Change

        output, _ = self.lstm(encoded_X, hidden_state)
        output, _ = pad_packed_sequence(output, batch_first=True)

        outs = []
        for i in range(X.size(0)):
            outs.append(output[i, seq_lengths[i] - 1,:])
        output = torch.stack(outs, dim=0)
        return self.output_layer(self.dropout_layer(output))
    
    def initalize_hidden_state(self, batch_size):
        return (

After Change

        seq_lengths, perm_idx = seq_lengths.sort(0, descending=True)
        X = X[perm_idx, :]
        inv_perm_idx = torch.tensor([i for i, _ in sorted(enumerate(perm_idx), key=lambda idx: idx[1])], dtype=torch.long)

        encoded_X = self.embedding(X)
        encoded_X = pack_padded_sequence(encoded_X, seq_lengths, batch_first=True)
        _, (ht, _) = self.lstm(encoded_X, hidden_state)

        return self.output_layer(self.dropout_layer(ht[-1][inv_perm_idx, :]))
    
    def initalize_hidden_state(self, batch_size):
        return (

Instances