def trim_hidden(inputs, hidden):
    """
    In the case where the last batch is smaller than batch_size,
    we will need to keep only the first N hidden states,
    where N is the number of samples in the last batch
    Args:
        inputs: the inputs in the last batch size
        hidden: the hidden state produced by the penultimate batch

    Returns: hidden
        the trimmed hidden state

    """
    batch_size = inputs.size(0)

    if isinstance(hidden, tuple):
        hidden_size = hidden[0].size(1)
    else:
        hidden_size = hidden.size(1)

    # do nothing
    if batch_size == hidden_size:
        return hidden

    # trim the hidden state to the remaining samples in the batch
    if isinstance(hidden, tuple):
        hidden = (hidden[0][:, :batch_size, :].contiguous(),
                  hidden[1][:, :batch_size, :].contiguous())
    else:
        hidden = hidden[:, :batch_size, :].contiguous()

    return hidden