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

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.

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?