Repository: Lightning-AI/torchmetrics

Issue:
MulticlassConfusionMatrix overflows for num_classes > 16
## 🐛 Bug

_multiclass_confusion_matrix_update() overflows for num_classes > 16, target.dtype=uint8 and preds.max()>15

### To Reproduce

```
import torch
from torchmetrics.classification import MulticlassConfusionMatrix

m = MulticlassConfusionMatrix(num_classes=17)
preds = torch.LongTensor([16])
targets = torch.ByteTensor([16])

m.update(preds, targets)
t = m.compute()
assert t[16, 16] == 1
assert t[1, 15] == 0
```

### Expected behavior

t[16, 16] == 1
t[1, 15] == 0

### Environment

- TorchMetrics version : fork master

### Additional context

Probably the same issue for **MulticlassStatScore**
```
def _multiclass_confusion_matrix_update(preds: Tensor, target: Tensor, num_classes: int) -> Tensor:
    """Compute the bins to update the confusion matrix with."""
    -- unique_mapping = (target * num_classes + preds).to(torch.long)
    ++ unique_mapping = target.to(torch.long) * num_classes + preds.to(torch.long)
    bins = _bincount(unique_mapping, minlength=num_classes**2)
    return bins.reshape(num_classes, num_classes)
```


Discussion:

Resolving patch:
diff --git a/src/torchmetrics/functional/classification/confusion_matrix.py b/src/torchmetrics/functional/classification/confusion_matrix.py
--- a/src/torchmetrics/functional/classification/confusion_matrix.py
+++ b/src/torchmetrics/functional/classification/confusion_matrix.py
@@ -322,7 +322,7 @@ def _multiclass_confusion_matrix_format(
 
 def _multiclass_confusion_matrix_update(preds: Tensor, target: Tensor, num_classes: int) -> Tensor:
     """Compute the bins to update the confusion matrix with."""
-    unique_mapping = (target * num_classes + preds).to(torch.long)
+    unique_mapping = target.to(torch.long) * num_classes + preds.to(torch.long)
     bins = _bincount(unique_mapping, minlength=num_classes**2)
     return bins.reshape(num_classes, num_classes)
 
diff --git a/src/torchmetrics/functional/classification/stat_scores.py b/src/torchmetrics/functional/classification/stat_scores.py
--- a/src/torchmetrics/functional/classification/stat_scores.py
+++ b/src/torchmetrics/functional/classification/stat_scores.py
@@ -399,7 +399,7 @@ def _multiclass_stat_scores_update(
             idx = target != ignore_index
             preds = preds[idx]
             target = target[idx]
-        unique_mapping = (target * num_classes + preds).to(torch.long)
+        unique_mapping = target.to(torch.long) * num_classes + preds.to(torch.long)
         bins = _bincount(unique_mapping, minlength=num_classes**2)
         confmat = bins.reshape(num_classes, num_classes)
         tp = confmat.diag()


Test patch:
diff --git a/tests/unittests/classification/test_confusion_matrix.py b/tests/unittests/classification/test_confusion_matrix.py
--- a/tests/unittests/classification/test_confusion_matrix.py
+++ b/tests/unittests/classification/test_confusion_matrix.py
@@ -218,6 +218,18 @@ def test_multiclass_confusion_matrix_dtype_gpu(self, input, dtype):
         )
 
 
+def test_multiclass_overflow():
+    """Test that multiclass computations does not overflow even on byte input."""
+    preds = torch.randint(20, (100,)).byte()
+    target = torch.randint(20, (100,)).byte()
+
+    m = MulticlassConfusionMatrix(num_classes=20)
+    res = m(preds, target)
+
+    compare = sk_confusion_matrix(target, preds)
+    assert torch.allclose(res, torch.tensor(compare))
+
+
 def _sklearn_confusion_matrix_multilabel(preds, target, normalize=None, ignore_index=None):
     preds = preds.numpy()
     target = target.numpy()
diff --git a/tests/unittests/classification/test_stat_scores.py b/tests/unittests/classification/test_stat_scores.py
--- a/tests/unittests/classification/test_stat_scores.py
+++ b/tests/unittests/classification/test_stat_scores.py
@@ -327,6 +327,24 @@ def test_top_k_multiclass(k, preds, target, average, expected):
     )
 
 
+def test_multiclass_overflow():
+    """Test that multiclass computations does not overflow even on byte input."""
+    preds = torch.randint(20, (100,)).byte()
+    target = torch.randint(20, (100,)).byte()
+
+    m = MulticlassStatScores(num_classes=20, average=None)
+    res = m(preds, target)
+
+    confmat = sk_confusion_matrix(target, preds)
+    fp = confmat.sum(axis=0) - np.diag(confmat)
+    fn = confmat.sum(axis=1) - np.diag(confmat)
+    tp = np.diag(confmat)
+    tn = confmat.sum() - (fp + fn + tp)
+    compare = np.stack([tp, fp, tn, fn, tp + fn]).T
+
+    assert torch.allclose(res, torch.tensor(compare))
+
+
 def _sklearn_stat_scores_multilabel(preds, target, ignore_index, multidim_average, average):
     preds = preds.numpy()
     target = target.numpy()