torch.nn.CrossEntropyLoss() torch.Size([32, 96, 128]) 3 dimension

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])

when chang it to Long type, or float type , it will get new error

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.

1 Like

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?