PyTorch how to do gathers over multiple dimensions

Question

I\'m trying to find a way to do this without for loops.

Say I have a multi-dimensional tensor t0:

bs = 4
seq = 10
v = 16
t0 = torch.rand((bs,