Thanks for the reply! How about when A has another dimension, now the shape of A is (batch size, hidden dim, data dim). And B is still the same, a long tensor with shape (batch size, data len). I actually want to have the resulting tensor with shape (batch size, data len, data dim).
For example, A=[[[5], [2], [6]], [[7], [3], [4]]] and B=[[0, 2, 1, 1], [2, 2, 1, 0]]. Then what I want is a tensor [[[5], [6], [2], [2]], [[4], [4], [3], [7]]]. Basically now data dim is 1 as an example.
It seems that A.gather(1, B) requires A and B to have the same shapes.