Mask tensors from indexes (e.g from multinomial sampling)

Hi,

I am looking for a ‘good’ way to transform matrix of indexes to a onehot matrix (line by line). For example,

Transforming[[1,3],[2,4]] to [[0,1,0,1,0],[0,0,1,0,1]]

Any idea ? Is there a way for example to extend index_fill_ to take into account 2D tensors instead of 1D tensor ?

1 Like

I’d recommend using scatter_:

>>> x = torch.LongTensor([[1, 3], [2, 4]])
>>> torch.zeros(2, 5).scatter_(1, x, 1)

 0  1  0  1  0
 0  0  1  0  1
[torch.FloatTensor of size 2x5]
2 Likes

how to reverse this step ?
i.e. from 2d mask tensor to 2d index ?

You could use this code snippet to get all indices of the value you are comparing to:

x = torch.LongTensor([[1, 3], [2, 4]])
res = torch.zeros(2, 5).scatter_(1, x, 1)
(res == 1.).nonzero()
1 Like