loss = torch.nn.CrossEntropyLoss()
loss_values = loss(train_output, train_label)
***** train_label shape: torch.Size([32, 96, 128])
***** train_output shape: torch.Size([32, 5, 96, 128])
I assume that 32 is your batch size. torch.nn.CrossEntropyLoss() is expecting you to pass output with dimensions [batch_size, n_classes, seq_len] and training label with dimensions [batch_size, seq_len] which cointains training class labels.
Also you can pass input with dimensions (N,C) and target with dimensions (N), from the documentation.
I do not know what your task is but I hope this will help you.
Hi, thanks. I want to use it for FCN pixels, now I don’t know how to check it. For FCN input and output have the same torch size. Do you have any idea?
What do 96 and 128 represent for size [32, 96, 128]?. Can you give more information about your task?
