Given a source tensor of shape
[B,N,F] and an index tensor of shape
index[i][j] is an index to a specific feature inside
source[i][j] is there a way to extract an output tensor such that:
output[i][j] = source[i][j][index[i][j]]
torch.gather specifies that
index.shape == source.shape, while here the shape of the source is one dimension bigger.
source = [ [[0.1,0.2],[0.2,0.3]], [[0.4,0.5],[0.6,0.7]], [[0.7,0.6],[0.8,0.9]] ] index = [ [1,0], [0,0], [1,1] ] desired_output = [ [0.2,0.2], [0.4,0.6], [0.6,0.9] ]