@albanD No surprise, your version is super fast! Using the following benchmark code, I obtained on my Mac (CPU):
10 trials took:
f1: 2.2150492668151855
f2: 0.36966371536254883
f3: 0.004826068878173828

Thank you so much again!
import time
import torch

def f1(x, idx1, idx2):
    """Using for loop
        x: N x B x V
        idx1: N x K matrix where idx1[i, j] is between [0, B)
        idx2: N x K matrix where idx2[i, j] is between [0, V)
    Return:
        y: N x K matrix where y[i, j] = x[i, idx1[i,j], idx2[i,j]]
    """
    # N x K
    y = torch.zeros(idx1.shape)
    N, K = idx1.shape
    for i in range(N):
        for j in range(K):
            y[i, j] = x[i, idx1[i,j], idx2[i,j]]
    return y

def f2(x, idx1, idx2):
    """Same as above but more efficient
    """
    # N x K
    return torch.einsum('iij->ij', x[:,idx1,idx2])


def f3(x, idx1, idx2):
    """Same as above, but even faster
    https://discuss.pytorch.org/t/similar-to-torch-gather-over-two-dimensions/118827/3?u=f10w
    """
    # Linearize the last two dims and index in a contiguous x
    x = x.contiguous()

    lin_idx = idx2 + x.size(-1) * idx1
    x = x.view(-1, x.size(1) * x.size(2))

    return x.gather(-1, lin_idx)


def main():
    N = 100
    B = 100
    V = 10000
    K = B*2

    repeat = 10
    t1 = 0
    t2 = 0
    t3 = 0

    for _ in range(repeat):
        x = torch.randn(N, B, V)
        idx1 = torch.randint(0, B, size=(N, K))
        idx2 = torch.randint(0, V, size=(N, K))

        start = time.time()
        y1 = f1(x, idx1, idx2)
        t1 += time.time() - start

        start = time.time()
        y2 = f2(x, idx1, idx2)
        t2 += time.time() - start

        start = time.time()
        y3 = f3(x, idx1, idx2)
        t3 += time.time() - start

        assert (y1-y2).abs().max() < 1e-10
        assert (y1-y3).abs().max() < 1e-10

    print(f'{repeat} trials took:\nf1: {t1}\nf2: {t2}\nf3: {t3}')

if __name__ == "__main__":
    main()