    def forward(
        self, inputs: torch.Tensor, states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        Given an input batch of sequential data such as word embeddings, produces
        a AugmentedLSTM representation of the sequential input and new state
        tensors.

        # Parameters

        inputs : `PackedSequence`, required.
            A tensor of shape (batch_size, num_timesteps, input_size)
            to apply the LSTM over.
        states : `Tuple[torch.Tensor, torch.Tensor]`
            Tuple of tensors containing
            the initial hidden state and the cell state of each element in
            the batch. Each of these tensors have a dimension of
            (bsize x num_layers x num_directions * nhid). Defaults to `None`.

        # Returns

        `Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]`
            AgumentedLSTM representation of input and
            the state of the LSTM `t = seq_len`.
            Shape of representation is (bsize x seq_len x representation_dim).
            Shape of each state is (bsize x num_layers * num_directions x nhid).

        """

        if not isinstance(inputs, PackedSequence):
            raise ConfigurationError("inputs must be PackedSequence but got %s" % (type(inputs)))

        # if states is not None:
        #    states = (states[0].transpose(0, 1), states[1].transpose(0, 1))
        if self.bidirectional:
            return self._forward_bidirectional(inputs, states)
        return self._forward_unidirectional(inputs, states)
Given an input batch of sequential data such as word embeddings, produces
a AugmentedLSTM representation of the sequential input and new state
tensors.

# Parameters

inputs : `PackedSequence`, required.
A tensor of shape (batch_size, num_timesteps, input_size)
to apply the LSTM over.
states : `Tuple[torch.Tensor, torch.Tensor]`
Tuple of tensors containing
the initial hidden state and the cell state of each element in
the batch. Each of these tensors have a dimension of
(bsize x num_layers x num_directions * nhid). Defaults to `None`.

# Returns

`Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]`
AgumentedLSTM representation of input and
the state of the LSTM `t = seq_len`.
Shape of representation is (bsize x seq_len x representation_dim).
Shape of each state is (bsize x num_layers * num_directions x nhid).
allennlp/modules/augmented_lstm.py