def index_select_ND(input, dim, index):
    """Wrapper around torch.index_select for non 1D index tensors."""
    input_shape = input.shape
    target_shape = index.shape + input_shape[dim + 1:]
    out_tensor = input.index_select(dim, index.view(-1))
    return out_tensor.view(target_shape)