Index Array to one-hot ID matrix

Question

Assume I have an array [0, 2], I would like to output a matrix consisting of one-hot vector based on [0, 2] like [ [1, 0, 0] [0, 0, 1]] (Note that the second dimension of output matrix is assumed to be 3 but it can be any number larger then argmax([0,2]) which is 2.

I can only think of this way achieves this function. Is there any simpler way.

t = torch.tensor([0,2])
dim2_size = 3
id_t = torch.zeros(t.shape[0], dim2_size)
row_idx = 0
for i in t:
  col_idx = i.item()
  id_t[row_idx, col_idx] = 1
  row_idx += 1
id_t

Answer 1

This one doesn't use any loop.

import torch

labels = torch.tensor([0, 2])
one_hot = torch.zeros(labels.shape[0], torch.max(labels)+1)
one_hot[torch.arange(labels.shape[0]), labels] = 1

print(one_hot)

tensor([[1., 0., 0.],
        [0., 0., 1.]])

Answer 2

method via numpy is more simple

import torch
import numpy as np

labels =[0,2]
output=np.eye(max(labels)+1)[labels]
print(torch.from_numpy(output))

Answer 3

In Pytorch this is best done through the use of scatter_ .

t = torch.tensor([0,2]).unsqueeze(0)
num_dims = 3
id_t = torch.zeros(num_dims, t.shape[1]).scatter_(0, t, 1)

This gives you id_t as:

tensor([[1., 0.],
        [0., 0.],
        [0., 1.]])