If the batch size can be derived from len(labels):

def to_onehot(labels, n_categories, dtype=torch.float32):
    batch_size = len(labels)
    one_hot_labels = torch.zeros(size=(batch_size, n_categories), dtype=dtype)
    for i, label in enumerate(labels):
        # Subtract 1 from each LongTensor because your
        # indexing starts at 1 and tensor indexing starts at 0
        label = torch.LongTensor(label) - 1
        one_hot_labels[i] = one_hot_labels[i].scatter_(dim=0, index=label, value=1.)
    return one_hot_labels

and you have 6 categories and want the output to be a tensor of integers:

to_onehot(labels, n_categories=6, dtype=torch.int64)

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

I would stick to torch.float32 in case you want to use label smoothing, mix-up or something along those lines later.

I try to find an optimized way without loop, but it seems like using loop is the only solution (to the best of my knowledge).

import torch

max_num = 4
multi_hot_num = 3
x = torch.tensor([0, 2, 1])

assert torch.all(x < max_num)

result = torch.zeros(len(x), max_num*multi_hot_num)
for i, (start, end) in enumerate(zip(x*multi_hot_num, (x+1)*multi_hot_num)):
    result[i, start:end] = 1.
    
print(result)

Result:

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

EDITED: without for loop

Solution based on this answer

import torch

max_num = 4
multi_hot_num = 3
x = torch.tensor([0, 2, 1])

assert torch.all(x < max_num)

start = (x * multi_hot_num).unsqueeze(-1)
end = ((x+1) * multi_hot_num).unsqueeze(-1)

result = torch.zeros(len(x), max_num*multi_hot_num)
index = torch.arange(max_num*multi_hot_num).unsqueeze(0).repeat(len(x), 1)

gte_start = start <= index
lt_end = index < end
mask = gte_start & lt_end
result[mask]=1

print(result)