I have a tensor `A`

of size `B x 4`

and another tensor `B`

of size `H x N`

. `A`

has integer values and many are -1 like:

```
tensor([[ 2.2640e+03, -1.0000e+00, -1.0000e+00, -1.0000e+00],
[ 3.5400e+02, -1.0000e+00, -1.0000e+00, -1.0000e+00],
[ 6.2700e+02, 6.2900e+02, -1.0000e+00, -1.0000e+00]
```

I want a new tensor of size `B x H`

with values from `B`

according to the indices in `A`

. For example, for the first entry, I have the tensor [2260, -1, -1, -1] and I want the tensor of size `H,`

in position `N`

(i.e., tensor2[:, 2260]). There might be entries from `A`

where there are more than one number higher than -1, in which case I would like to get the respective tensors in the indices and then average them.

`torch.index_select(b, 0, a[:, 0])`

almost gives the correct answer. Just need a way to handle the case where more than 1 integer is given, like in the third row of the above example.

Thanks for any help you can provide.