PyTorch中二分类0/1张量调用one_hot编码报错如何调试

报错原因 #

F.one_hot 仅支持整数类型的索引张量作为输入，你生成 capital_t 时使用的 torch.Tensor(y) 会默认生成 torch.float32 类型的浮点张量，不符合接口参数要求，因此触发报错。

解决方案 #

提供三个可直接落地的修改方案：

• 方案1：生成张量时直接指定整数类型
  将 capital_t = torch.Tensor(y) 替换为如下代码，直接生成符合要求的长整型张量：

  capital_t = torch.tensor(y, dtype=torch.long)
  

• 方案2：调用one_hot前做类型转换
  若不想修改张量生成逻辑，调用one_hot时追加类型转换逻辑即可：

  one_hot_capital = F.one_hot(capital_t.to(torch.long), 2)
  

• 方案3：用torch.as_tensor生成指定类型的张量
  对于列表输入也可以用as_tensor方法指定类型，内存占用更高效：

  capital_t = torch.as_tensor(y, dtype=torch.int64)
  

验证方法 #

修改后可以先输出张量类型确认是否符合要求：

print(capital_t.dtype)
# 正确输出应为 torch.int64 或 torch.long

内容的提问来源于stack exchange，提问作者riskiem