Using Focal Loss, error with num_classes

[Likhith-Sugganahalli]

definition of the loss in the model config:

loss function of the base semantic head:

outputs of the debug statements in the image above:
seg_preds torch.Size([2, 2, 192, 336])
gt_semantic_seg torch.Size([2, 1, 192, 336])
temp1 torch.Size([129024, 2])
temp2 torch.Size([129024])

definition of the sigmoid_focal_loss function

output of the debug statements:
num_classes for focal loss: 2
pred torch.Size([129024, 2])
target torch.Size([129024])

loss = _sigmoid_focal_loss( File "/home/ec2-user/.local/lib/python3.9/site-packages/mmcv/ops/focal_loss.py", line 60, in forward ext_module.sigmoid_focal_loss_forward( RuntimeError: target.max().item<int64_t>() <= (int64_t)num_classes INTERNAL ASSERT FAILED at "/tmp/mmcv/mmcv/ops/csrc/pytorch/cuda/focal_loss_cuda.cu":12, please report a bug to PyTorch. target label should smaller or equal than num classes