About slicing a tensor by an index tensor

1

This question may be basic and there may be a simple good way,
but I took not short time for this.

I want to silice tensor by another tensor.
Suppose we have a tensor of size [num_grid, num_box, len(x, y, w, h)]
in the situation that all grids has some boxes and its coodinates,
and a tensor of size [num_grid] which means that each index is the maximum IoU box for each grid.

Is there a good way to obtain a tensor of size [num_grid, len(x, y, w, h)]
the tensor of the highest IoU box for each grid?
Thank you.

2

This code should work:

num_grid, num_box, c = 3, 4, 5
x = torch.randn(num_grid, num_box, c)
idx = torch.randint(0, num_box, (num_grid,))
x[torch.arange(num_grid), idx]
3

Sorry for my late response.
The code is what I’ve wanted and looks a very simple&good solution!
Thank you for the helpful advice!