x is (B, T, D) of right-padded sequences and lengths is (B,). Return the (B, D) embedding at each sequence's last real token, which sits at index length - 1.

Do not loop over the batch. Taking x[:, -1] instead would hand you padding for every sequence shorter than the longest -- the mistake this problem is about.

Input

x =
tensor([[[ 0.,  1.,  2.],
         [ 3.,  4.,  5.],
         [ 6.,  7.,  8.],
         [ 9., 10., 11.]],

        [[12., 13., 14.],
         [15., 16., 17.],
         [18., 19., 20.],
         [21., 22., 23.]]])
lengths = tensor([2, 4])

Output

tensor([[ 3.,  4.,  5.],
        [21., 22., 23.]])