import torch


def select_embs(embs, pred_cands, M):
    """Select the embeddings for the predicted indices.

    Args:
        embs: embedding tensor (M x K x H)
        pred_cands: indices into the K candidates to select (M)
        M: M

    Returns: selected embedding (M x H)
    """
    return embs[torch.arange(M).unsqueeze(0), pred_cands]