TonyTr:

I need to compute the loss on the first dim,


I am not sure I understand this. If your labels are a [3,4] tensor, what does this translate to? You need to compute the loss on the first row, on the first column, or something else?



 TonyTr:

How do I test the loss function against toy inputs?


You could use random arrays. I did the following to test the loss function from your code above:
x = torch.rand(3, 4, 5)
torch.sum(-x * F.log_softmax(x, dim=1), dim=1)