I don’t think there’s a native way to do this. You could do this using the for comprehension though:
xs = [2, 1, 2, 0]
ys = [0, 1, 2, 3]
torch.cat([a[x, :, y, :].unsqueeze(0) for x, y in zip(xs, ys)])
1 Like