Assume I have a model that takes as input a sequence of vectors i.e. [B,N,C] = input.shape with B as batch dimension, N as sequence length and C as latent dim.
Additionaly at some point in my model I get a tensor w, that is of
shape [B,N] that weights each element of my sequence i.e. w.sum(dim=1) == [1., 1., 1., ...] of size [B,]
What the follwing code showes is my attempt to get the top k sequence elements of my input X that have the highest values in w:
X = torch.rand(4,7,128)
w = torch.rand(4,7)
w_sorted = torch.argsort(w, dim = 1, descending = True)
w_trunc = w_sorted[:,:5]

X_topk = torch.empty(X.shape[0], 5, 128)
for i in range(X.shape[0]):
    X_topk[i,:,:] = X[i,w_trunc[i,:],:]  

Is there is a better way to do this without the loop and directly indexing multiple dimensions of X using w_trunc?