I’m not sure to fully understand your use case, however it seems you are dealing with a two-class classification (2 output channels and two class indices).
If your output has the shape [B, 2, H, W], your target should have the shape [B, H, W] and contain the class indices in the range [0, 1] to be able to use nn.CrossEntropyLoss.
Also, make sure to pass raw logits to this criterion, as internally F.log_softmax and nn.NLLLoss will be applied. If you are using a softmax activation as the last non-linearity, could you please remove it?
Could you check your inputs for invalid values (Inf or NaN)?
If you can’t find any invalid values, could you lower the learning rate a bit and check, if you get a NaN loss again?