How to get subset tensor by a tensor of indices (Pytorch)

Given:

• a tensor A of has shape: (batch_size, vocab_size)
• a tensor B is tensor of token indices has shape: (batch_size, seq_len

I want to get the values of A at indices B.For example:

A = [[ 0.6301,  0.2310, -1.1964, -1.1293, -1.0428,  0.4011,  0.0519, -3.0591],     [ 1.3691, -0.7477, -0.9323, -0.3670, -0.1568, -0.1282, -0.4557,  0.1747]])B = [[0, 2, 7, 7, 5],     [2, 4, 7, 2, 5]])output = [[ 0.6301, -1.1964, -3.0591, -3.0591,  0.4011],          [-0.9323, -0.1568,  0.1747, -0.9323, -0.1282]]

Thanks you