I have a 2D tensor A of size (n, d) where each row has d elements
and another 2D tensor mask of size (n, 10) where each row has 10 integers in range [0, d-1].

I wish to extract elements from the (n,d) tensor using indices of (n,10), to obtain (n, 10) elements from A.