Direct indexing should work:
a = torch.tensor([[[ 0,  1,  2],
                   [ 3,  4,  5],
                   [ 6,  7,  8],
                   [ 9, 10, 11]],

                  [[12, 13, 14],
                   [15, 16, 17],
                   [18, 19, 20],
                   [21, 22, 23]]]).float()

b = torch.tensor([[0, 1], [1, 2]])

print(a[b[0, :], b[1, :]])
> tensor([[ 3.,  4.,  5.],
          [18., 19., 20.]])