PyTorch：将张量映射为最大值对应位置的One-Hot张量的简洁实现

错误原因 #

torch.nn.functional.one_hot要求输入必须是整数类型的索引张量，你直接传入原始的浮点型数值张量，自然会触发RuntimeError。

一行代码实现方案 #

用纯PyTorch原生API实现，先通过argmax获取最大值对应的索引（自动转为整数类型），再传入one_hot，最后转成浮点型（默认one_hot输出为长整型）：

b = torch.nn.functional.one_hot(a.argmax(), num_classes=4).float()

验证效果 #

针对输入张量a = torch.Tensor([0.0917, -0.0006, 0.1825, -0.2484])，执行上述代码后，b的输出结果为：

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

额外说明 #

该方案无需依赖numpy，完全基于PyTorch实现，比你原来的代码更简洁且符合PyTorch的原生使用习惯。

内容的提问来源于stack exchange，提问作者Mihai.Mehe