Pytorch tensor转onehot

def test_onehot():
    v = torch.tensor([[0.1, 0.2, 0.7],
                      [0.1, 0.6, 0.3],
                      [0.1, 0.5, 0.4],
                      [0.8, 0.1, 0.1], ])

    print('v', v.size(), v)
    # 按照形状创建全0张量
    result = torch.zeros_like(v, dtype=torch.long)
    # 目标维度
    dim = -1
    # 根据索引将值改为1
    result.scatter_(dim,
                    v.argmax(dim).unsqueeze(dim),
                    torch.ones(4, dtype=torch.long).unsqueeze(dim))

    print('result', result.size(), result)

v torch.Size([4, 3]) tensor([[0.1000, 0.2000, 0.7000],
        [0.1000, 0.6000, 0.3000],
        [0.1000, 0.5000, 0.4000],
        [0.8000, 0.1000, 0.1000]])
result torch.Size([4, 3]) tensor([[0, 0, 1],
        [0, 1, 0],
        [0, 1, 0],
        [1, 0, 0]])