我有一个基于BERT的序列分类模型，它将4个字符串作为输入，并为每个字符串输出2个标签：

my_input = [string_1, string_2, string_3, string_4]
out_logits = model(my_input).logits
out_softmax = torch.softmax(out_logits)
out_softmax 
>>> tensor([[0.8666, 0.1334],
        [0.8686, 0.1314],
        [0.8673, 0.1327],
        [0.8665, 0.1335]], device='cuda:0', grad_fn=<SoftmaxBackward0>)

我的损失函数是nn.CrossEntropyLoss()，我的标签是Tensor，其索引对应于正确的标签：tensor([0., 0., 0., 1.])。注意，除了一个标签之外，每个标签都是1。

loss = loss_fun(out_softmax, labels_tensor)
# step
optim.zero_grad()
loss.backward()
optim.step()

我遇到的问题是，模型只能预测一个类（例如，上面的第一列），不完全确定为什么会发生这种情况，但我认为对应该是1的预测进行更多的惩罚可能会有所帮助。
我怎样才能惩罚更多的预测？