使用带权重的CrossEntropyLoss报错：期望Float类型却得到Long

我在用前馈神经网络执行4分类任务，因类别不平衡计划给CrossEntropyLoss添加权重，但设置后出现RuntimeError: expected scalar type Float but found Long错误。

原权重设置代码：

loss_fn = nn.CrossEntropyLoss(weight=torch.tensor([1, 1000, 1000, 1000]))

错误原因 #

PyTorch的CrossEntropyLoss要求权重张量的数据类型为浮点型（Float），但直接创建的torch.tensor([1,1000,1000,1000])默认是长整型（Long），和模型输出的浮点型张量类型不匹配，触发类型错误。

解决方案 #

有两种方式修正权重张量的类型：

1. 创建时指定浮点类型

# 显式指定dtype
loss_fn = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 1000.0, 1000.0, 1000.0], dtype=torch.float))
# 或用浮点数值初始化，自动推导类型
loss_fn = nn.CrossEntropyLoss(weight=torch.tensor([1., 1000., 1000., 1000.]))

2. 对已有张量做类型转换

weights = torch.tensor([1, 1000, 1000, 1000]).float()
loss_fn = nn.CrossEntropyLoss(weight=weights)

额外注意 #

若模型运行在GPU上，需将权重同步到对应设备：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
weights = torch.tensor([1., 1000., 1000., 1000.]).to(device)
loss_fn = nn.CrossEntropyLoss(weight=weights)

修改后的完整示例 #

ffn_model = FeedforwardNeuralNetModel(input_dim=32, hidden_dim=128, output_dim=4)
# 带浮点权重的损失函数
loss_fn = nn.CrossEntropyLoss(weight=torch.tensor([1., 1000., 1000., 1000.]))
optimizer = torch.optim.SGD(ffn_model.parameters(), lr=0.1)

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