How to effectively construct a mask tensor?

1

Suppose I have a tensor indicating which column should be 1 for each row, for example,

index = torch.tensor([3,1,0,0,2])

and I would like to construct a mask tensor from above one and get this result:

mask = torch.tensor([[0,0,0,1,0],
[0,1,0,0,0],
[1,0,0,0,0],
[1,0,0,0,0],
[0,0,1,0,0]])

Is there any pytorch api can help me finish this job? Or maybe with several apis?

2

scatter_ should work:

index = torch.tensor([3,1,0,0,2])
torch.zeros(index.size(0), 5).scatter_(1, index.unsqueeze(1), 1.)
# tensor([[0., 0., 0., 1., 0.],
#         [0., 1., 0., 0., 0.],
#         [1., 0., 0., 0., 0.],
#         [1., 0., 0., 0., 0.],
#         [0., 0., 1., 0., 0.]])

I don’t know how the mask shape was defined in dim1 so I just used 5 as seen in your example.

3

Wow, this works perfectly! Thank you!