import torch


def get_last_token(t, lengths):
    """
    Grab last hidden activation of each batch element according to `lengths`

    #                               ^ (1)      ^ (2)      ^ (3)
    >>> t = torch.tensor([[[1],[2],[3]], [[2],[3],[4]], [[3],[4],[5]]])
    >>> lengths = torch.tensor([3, 2, 1])
    >>> get_last_token(t, lengths).tolist()
    [[3], [3], [3]]
    """
    if not isinstance(lengths, torch.Tensor):
        lengths = torch.tensor(lengths)

    seq_len, batch, _ = t.size()
    index = torch.arange(0, batch, dtype=torch.int64, device=t.device) * seq_len
    index = index + (lengths - 1)
    t = t.transpose(0, 1).contiguous()  # make it batch first
    t = t.view(seq_len * batch, -1)
    t = t.index_select(0, index)
    t = t.view(batch, -1)
    return t