RuntimeError: The size of tensor a (128) must match the size of tensor b (2) at non-singleton dimension 4

1

I am using BCEWithLogitsLoss as loss function, and the output size and targets size is the same as ·torch.Size([4, 1, 128, 128, 128])· but I run into that error.

2

nn.BCEWithLogitsLoss should work, if the model output and target have the same shape.
Could you print the shape of both inputs again and check, where the size of 2 is coming from?