hey, I have a similar question regarding cross entropy, but for multi-class prediction. My output of network, the logits has shape of [5,6,256,256].  Number of classes is 6 here. The target has a shape of [5,256,256]. Since this criterion combines  LogSoftMax  and ClassNLLCriterion in one single class, cross entropy expects logits and target having different size, right? At least,  criterion = nn.CrossEntropyLoss()  loss = criterion(logit, true_masks) didn’t give me error.
And, I have another question regarding to the loss value of cross entropy. My initial values are larger than 1. According to some sanity checks post out there, it says that it is about bad initialization of weights. My activation function is leaky_relu. So, I used He distribution to initialize the weights. For example like this,
        self.down = nn.Sequential(             nn.ZeroPad2d(1),             nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=scale, stride=scale))         self.down[1].weight = nn.Parameter(             nn.init.kaiming_normal_(torch.empty((out_channels, in_channels,scale,scale),dtype=torch.float32),                                     nonlinearity='leaky_relu'),requires_grad = True)
I compared to the case of using the default initialization. Using He distribution gives me even larger initial losses. My validation accuracy is lower than 50% for first 30 epochs which I have to keep exploring, since the data is complex data. Can you please give me any hints? Thank u.