def gather_sequence(input_sequence, indices):
    """
    Given a batch of sequence, gather an element from each sequence
    :param input_sequence: [B, T, ...]
    :param indices: [T, ...]
    :return: [B, ...]
    """
    # make inds shape [1, T, ...]
    indices = indices.view(-1, *([1]*(input_sequence.ndimension() - 1)))
    out = input_sequence.gather(1, indices.expand(-1, 1, *input_sequence.shape[2:]))
    return out.squeeze(1)