Thanks, it appears tensor.gather(...) is the right tool. Unfortunately, I still couldn’t get it to work. When I add an extra dimension at the end of adj and use keepdim=True, features.gather(1,m) only returns a tensor of size 4 x 1 x 1 while I need the whole 4 x 10 x 1 tensor. So gather somewhat ignores the last dimension of size 10 of features, it only picks the first element of that dimension.
Update: So I got it to work by using m.repeat(1,1,10) and using that as index tensor for the gather call. This still doesn’t seem like the most elegant solution to me. Now I gotta figure out how to do this with batches.
Update 2: So batching it is no problem, it’s literally the same with an additional dimension. I’m still interested whether there is a better way than calling repeat.