Getting all tensor elements but elements from specific indices

Hi, I’m having the following problem which I can’t seem to resolve easily.

I have the following 2 tensors ‘values’ & ‘idx’ where:

values.shape = [2048,16,4]
idx.shape = [2048,16,1]

the elements in the tensor ‘idx’ are ranging from [0,1,2,3], and they indicate a specific elements in the corresponding place in the ‘values’ vector.
I wan’t to be able to separate the tensor ‘values’ according to the ‘idx’ tensor such that:

values_a.shape = [2048,16,1]
values_b.shape = [2048,16,3]

where the elements in ‘values_a’ are chosen according to the ‘idx’ tensor
and ‘values_b’ is chosen by the elements not indicated in the ‘idx’ tensor

I managed to get ‘values_a’ using gather function.
But couldn’t figure a way to get the other.
Any help please?

idx_all = torch.arange(4).view(1,1,4).expand_as(values)
values_b = values[mask].view(2048,16,3)

No strict check yet, but it should work in similar way.

second line shows this error:

RuntimeError: shape '[2048, 16, 4]' is invalid for input of size 32768

Sorry it should be mask=idx.expand_as(idx_all)!=idx_all

1 Like