Applies one step of an lstm (plus reduction) to the input X, which
        is handled by self.encoder

def forward(self, X):
        """Applies one step of an lstm (plus reduction) to the input X, which
        is handled by self.encoder"""
        # Identify the first non-zero integer from the right (i.e., the length
        # of the sequence before padding starts).
        batch_size, max_seq = X.shape[0], X.shape[1]
        seq_lengths = torch.zeros(batch_size, dtype=torch.long)
        for i in range(batch_size):
            for j in range(max_seq - 1, -1, -1):
                if not torch.all(X[i, j] == 0):
                    seq_lengths[i] = j + 1
                    break

        # Sort by length because pack_padded_sequence requires it
        # Save original order to restore before returning
        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,
        )

        # Encode and pack input sequence
        X_packed = rnn_utils.pack_padded_sequence(
            self.encoder.encode(X), seq_lengths, batch_first=True
        )

        # Run LSTM
        outputs, (h_t, c_t) = self.lstm(X_packed)

        # Unpack and reduce outputs
        outputs_unpacked, _ = rnn_utils.pad_packed_sequence(outputs, batch_first=True)
        reduced = self._reduce_output(outputs_unpacked, seq_lengths)
        return reduced[inv_perm_idx, :]