[BUG] BUG in cross_entropy_loss.py

[Dawn-bin]

Describe the bug
hi, I copy code to use in my own project, and I find a Issuse when I use cross_entropy_loss.
In line 121

if pred.size(1) == 1:
        # For binary class segmentation, the shape of pred is
        # [N, 1, H, W] and that of label is [N, H, W].
        assert label.max() <= 1, \
            'For pred with shape [N, 1, H, W], its label must have at ' \
            'most 2 classes'
        pred = pred.squeeze()

Should ' label.max() <= 1' mask out ignore_index? Since the ignore_index often set as 255.

Bug fix

if pred.size(1) == 1:
        # For binary class segmentation, the shape of pred is
        # [N, 1, H, W] and that of label is [N, H, W].
        assert label[label != ignore_index].max() <= 1, \
            'For pred with shape [N, 1, H, W], its label must have at ' \
            'most 2 classes'
        pred = pred.squeeze()

[MengzhangLI]

Hi, thanks for your suggestion. We think you are correct, could you make a PR to fix this problem? Thanks in advance.

[marz869]

Hi there, I faced another problem. In the case that labels are float32, label.max()<=1 returns false.
for example:
tensor(1.0000)<=1 returns false. That might be due to the internal representation of tensor(1.0000) which is not precisely equal to 1,
to solve the problem I change this line:
assert label.max() <= 1, "For pred with shape [N, 1, H, W], its label must have at " "most 2 classes"
to this line:
assert torch.allclose(label.max(), torch.tensor(1, dtype=torch.float, device=label.device)), "For pred with shape [N, 1, H, W], its label must have at " "most 2 classes"