Loss function error

1

when i use nn.CrossEntropyLoss to train a segmentation model.But it shows that 'lonly batches of spatial targets supported(non-empty 3D tensors)but got targets of size :[2,1,512,512]
the output_size of the model is [2,2,512,512](batch_size = 2 and num_classes = 2]

Thanks

2

Have you tried squeezeing the targets to [2, 512, 512]?

1 Like
3

Thank you very much. When i squeeze ing the targets to [2, 512, 512] ,it is work.

1 Like