Repository: Lightning-AI/torchmetrics

Issue:
ConfusionMatrix is float32 which gives slightly varying ground truth counts
## 🐛 Bug

confmat in ConfusionMatrix is of data type float32 (see: https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/classification/confusion_matrix.py#L121). This results in slightly varying ground truth counts when the confusion matrix contains a lot of values (e.g. for semantic segmentation):

### To Reproduce

Steps to reproduce the behavior:

```python
import torch
from torchmetrics import ConfusionMatrix

batchsize = 256
n_classes = 2
width = 512
height = 512

metric = ConfusionMatrix(num_classes=n_classes)

target = torch.randint(size=(batchsize, height, width), low=0, high=n_classes)

for i in range(10):
    preds = torch.randint(size=(batchsize, height, width), low=0, high=n_classes)
    metric.update(preds=preds, target=target)
    print(torch.sum(metric.confmat, axis=1))
    metric.reset()
```

The output is:

```bash
tensor([33556800., 33552060.])
tensor([33556804., 33552060.])
tensor([33556800., 33552060.])
tensor([33556804., 33552060.])
tensor([33556800., 33552060.])
tensor([33556804., 33552060.])
tensor([33556804., 33552060.])
tensor([33556804., 33552060.])
tensor([33556800., 33552060.])
tensor([33556804., 33552060.])
```

### Expected behavior

The number of ground truth pixels should not vary. When I change the data type of confmat to int, everything is fine.

### Environment

- PyTorch Version (e.g., 1.0): 1.10.0
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): conda
- Python version: 3.8


Discussion:
Hi! thanks for your contribution!, great first issue!
interesting, probably using float64 or uint64 would help... are you interested in sending fix? :rabbit: 

Resolving patch:
diff --git a/torchmetrics/classification/confusion_matrix.py b/torchmetrics/classification/confusion_matrix.py
--- a/torchmetrics/classification/confusion_matrix.py
+++ b/torchmetrics/classification/confusion_matrix.py
@@ -69,26 +69,26 @@ class ConfusionMatrix(Metric):
         >>> preds = torch.tensor([0, 1, 0, 0])
         >>> confmat = ConfusionMatrix(num_classes=2)
         >>> confmat(preds, target)
-        tensor([[2., 0.],
-                [1., 1.]])
+        tensor([[2, 0],
+                [1, 1]])
 
     Example (multiclass data):
         >>> target = torch.tensor([2, 1, 0, 0])
         >>> preds = torch.tensor([2, 1, 0, 1])
         >>> confmat = ConfusionMatrix(num_classes=3)
         >>> confmat(preds, target)
-        tensor([[1., 1., 0.],
-                [0., 1., 0.],
-                [0., 0., 1.]])
+        tensor([[1, 1, 0],
+                [0, 1, 0],
+                [0, 0, 1]])
 
     Example (multilabel data):
         >>> target = torch.tensor([[0, 1, 0], [1, 0, 1]])
         >>> preds = torch.tensor([[0, 0, 1], [1, 0, 1]])
         >>> confmat = ConfusionMatrix(num_classes=3, multilabel=True)
         >>> confmat(preds, target)  # doctest: +NORMALIZE_WHITESPACE
-        tensor([[[1., 0.], [0., 1.]],
-                [[1., 0.], [1., 0.]],
-                [[0., 1.], [0., 1.]]])
+        tensor([[[1, 0], [0, 1]],
+                [[1, 0], [1, 0]],
+                [[0, 1], [0, 1]]])
 
     """
     is_differentiable = False
@@ -118,7 +118,10 @@ def __init__(
         if self.normalize not in allowed_normalize:
             raise ValueError(f"Argument average needs to one of the following: {allowed_normalize}")
 
-        default = torch.zeros(num_classes, 2, 2) if multilabel else torch.zeros(num_classes, num_classes)
+        if multilabel:
+            default = torch.zeros(num_classes, 2, 2, dtype=torch.long)
+        else:
+            default = torch.zeros(num_classes, num_classes, dtype=torch.long)
         self.add_state("confmat", default=default, dist_reduce_fx="sum")
 
     def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
diff --git a/torchmetrics/functional/classification/confusion_matrix.py b/torchmetrics/functional/classification/confusion_matrix.py
--- a/torchmetrics/functional/classification/confusion_matrix.py
+++ b/torchmetrics/functional/classification/confusion_matrix.py
@@ -158,26 +158,26 @@ def confusion_matrix(
         >>> preds = torch.tensor([0, 1, 0, 0])
         >>> confmat = ConfusionMatrix(num_classes=2)
         >>> confmat(preds, target)
-        tensor([[2., 0.],
-                [1., 1.]])
+        tensor([[2, 0],
+                [1, 1]])
 
     Example (multiclass data):
         >>> target = torch.tensor([2, 1, 0, 0])
         >>> preds = torch.tensor([2, 1, 0, 1])
         >>> confmat = ConfusionMatrix(num_classes=3)
         >>> confmat(preds, target)
-        tensor([[1., 1., 0.],
-                [0., 1., 0.],
-                [0., 0., 1.]])
+        tensor([[1, 1, 0],
+                [0, 1, 0],
+                [0, 0, 1]])
 
     Example (multilabel data):
         >>> target = torch.tensor([[0, 1, 0], [1, 0, 1]])
         >>> preds = torch.tensor([[0, 0, 1], [1, 0, 1]])
         >>> confmat = ConfusionMatrix(num_classes=3, multilabel=True)
         >>> confmat(preds, target)  # doctest: +NORMALIZE_WHITESPACE
-        tensor([[[1., 0.], [0., 1.]],
-                [[1., 0.], [1., 0.]],
-                [[0., 1.], [0., 1.]]])
+        tensor([[[1, 0], [0, 1]],
+                [[1, 0], [1, 0]],
+                [[0, 1], [0, 1]]])
 
     """
     confmat = _confusion_matrix_update(preds, target, num_classes, threshold, multilabel)


Test patch: