Repository: Lightning-AI/torchmetrics

Issue:
Explicitly initializing tensors in `float32` in MeanMetric goes against `torch.set_default_dtype`, leading to numerical errors
## 🐛 Bug

When working in `torch.float64` (set by `torch.set_default_dtype(torch.float64)`), the result of `MeanMetric` is slightly incorrect because, internaly, the `weight` and `value` tensors are initialized to `torch.float32` explicitly (see these [lines](https://github.com/Lightning-AI/torchmetrics/blob/c0d2d3ab7221f8804ba55baf3945637d68485cc1/src/torchmetrics/aggregation.py#L558)).

Naturally, if someone's working in float64 then the problem can't tolerate the slight mismatch.

### To Reproduce

```python
import torch
from torchmetrics.aggregation import MeanMetric

torch.set_default_dtype(torch.float64)

metric = MeanMetric()

values = torch.randn(10000)
metric.update(values)
result = metric.compute()

print(f"{result} = Result\n{values.mean()} = Actual mean")
```

The output of the snippet is:
```
0.018805044555664063 = Result
0.01880504383120193 = Actual mean
```

Notice that the two differ. Where as if we switch the default to `torch.float32`:
```
0.013675613328814507 = Result
0.013675613328814507 = Actual mean
```
they're exactly the same.

### Expected behavior

Idealy, the result of the MeanMetric (and possibly other metrics where this might have happened) should give the correct result irrespective of the default dtype.

### Environment

`pip freeze`:

```
torch==2.1.0+cu121
torchmetrics==1.2.1
```

Python 3.9

Windows 11



Discussion:
Hi! thanks for your contribution!, great first issue!
So honoring the default if it is set would help, right? mind sending a fix PR... :rabbit: 
Something like that.

I could try, would be my first time but I'll give it a go.

Resolving patch:
diff --git a/src/torchmetrics/aggregation.py b/src/torchmetrics/aggregation.py
--- a/src/torchmetrics/aggregation.py
+++ b/src/torchmetrics/aggregation.py
@@ -157,7 +157,7 @@ def __init__(
     ) -> None:
         super().__init__(
             "max",
-            -torch.tensor(float("inf")),
+            -torch.tensor(float("inf"), dtype=torch.get_default_dtype()),
             nan_strategy,
             state_name="max_value",
             **kwargs,
@@ -262,7 +262,7 @@ def __init__(
     ) -> None:
         super().__init__(
             "min",
-            torch.tensor(float("inf")),
+            torch.tensor(float("inf"), dtype=torch.get_default_dtype()),
             nan_strategy,
             state_name="min_value",
             **kwargs,
@@ -366,7 +366,7 @@ def __init__(
     ) -> None:
         super().__init__(
             "sum",
-            torch.tensor(0.0),
+            torch.tensor(0.0, dtype=torch.get_default_dtype()),
             nan_strategy,
             state_name="sum_value",
             **kwargs,
@@ -536,12 +536,12 @@ def __init__(
     ) -> None:
         super().__init__(
             "sum",
-            torch.tensor(0.0),
+            torch.tensor(0.0, dtype=torch.get_default_dtype()),
             nan_strategy,
             state_name="mean_value",
             **kwargs,
         )
-        self.add_state("weight", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("weight", default=torch.tensor(0.0, dtype=torch.get_default_dtype()), dist_reduce_fx="sum")
 
     def update(self, value: Union[float, Tensor], weight: Union[float, Tensor] = 1.0) -> None:
         """Update state with data.


Test patch:
diff --git a/tests/unittests/bases/test_aggregation.py b/tests/unittests/bases/test_aggregation.py
--- a/tests/unittests/bases/test_aggregation.py
+++ b/tests/unittests/bases/test_aggregation.py
@@ -204,3 +204,20 @@ def test_mean_metric_broadcast(nan_strategy):
     metric.update(x, w)
     res = metric.compute()
     assert round(res.item(), 4) == 3.2222  # (0*0 + 2*2 + 3*3 + 4*4) / (0 + 2 + 3 + 4)
+
+
+@pytest.mark.parametrize(
+    ("metric_class", "compare_function"),
+    [(MinMetric, torch.min), (MaxMetric, torch.max), (SumMetric, torch.sum), (MeanMetric, torch.mean)],
+)
+def test_with_default_dtype(metric_class, compare_function):
+    """Test that the metric works with a default dtype of float64."""
+    torch.set_default_dtype(torch.float64)
+    metric = metric_class()
+    values = torch.randn(10000)
+    metric.update(values)
+    result = metric.compute()
+    assert result.dtype == torch.float64
+    assert result.dtype == values.dtype
+    assert torch.allclose(result, compare_function(values), atol=1e-12)
+    torch.set_default_dtype(torch.float32)