Pytorch transform tensor to one hot
What is the easiest way to transform tensor of shape (batch_size, height, width) filled with n values to tensor of shape (batch_size, n, height, width)? 
I created solution below, but looks like there are easier and faster way to do this


def batch_tensor_to_onehot(tnsr, classes):
    tnsr = tnsr.unsqueeze(1)
    res = []
    for cls in range(classes):
        res.append((tnsr == cls).long())
    return torch.cat(res, dim=1)


You can use torch.nn.functional.one_hot.

For your case:

a = torch.nn.functional.one_hot(tnsr, num_classes=classes)
out = a.permute(0, 3, 1, 2)

