Applying cross entropy loss on one-hot targets

Hi, I have labels in one-hot format with size [bsz, bsz2]. My input also is a matrix of shape [bsz,bsz2]. I want to use cross-entropy loss. I searched the pytorch doc and I found that we can’t apply cross-entropy loss on one hot except in the following way:

out = torch.FloatTensor([[0.05, 0.9, 0.05], [0.05, 0.05, 0.9], [0.9, 0.05, 0.05]])

y1 = torch.FloatTensor([[0, 1, 0], [0, 0, 1], [1, 0, 0]])

_, targets = y1.max(dim=0)
loss = nn.CrossEntropyLoss()(out, Variable(targets))
print(loss.item())

################################

a = torch.randn(3,3)
print(a.shape)

b = torch.randn(3,3)

c = torch.cat((a,b),dim=1)

lable = torch.randn(3,3)

_,target =lable.max(dim=0)

criterion = torch.nn.CrossEntropyLoss(reduction="sum")

loss = criterion(Variable(target),c)

IndexError: Dimension out of range (expected to be in the range of [-1, 0], but got 1)

I can’t understand why I get an error in the second case. Any help is really appreciated.

Try to apply the max operation on dim1:

target = torch.argmax(lable, 1)

as dim0 is the batch dimension, while dim1 should correspond to the “class dimension”.