Repository: Lightning-AI/torchmetrics

Issue:
PyTorch Lightning's Precision 64 and version 0.8: Expected scalar type double
## 🐛 Bug

Since v0.8 (honestly, did only try 0.8.1 and 0.8.2) setting `precision` to 64 will crash any `stat_scores` dependent metric.

### To Reproduce

Just use any `stat_scores` metric within any trainer call, be it `validate` or `fit`.

#### Code sample
```python
import torch
from pytorch_lightning import LightningModule, Trainer
from torch.utils.data import TensorDataset, DataLoader
from torchmetrics import Accuracy


class LitModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.metric = Accuracy()

    def validation_step(self, *args, **kwargs):
        self.metric(torch.tensor([50, 50, 50, 50]), torch.tensor([11, 9, 14, 22]))


def dummy_dataloader():
    inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
    tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
    return DataLoader(TensorDataset(inps, tgts), batch_size=2)


if __name__ == '__main__':
    dataloader = dummy_dataloader()

    model = LitModel()
    trainer = Trainer(precision=64)
    trainer.validate(model, dataloader)
```

### Expected behavior

It shouldn't crash and calculate metrics as in any other precision.

### Environment

- pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel
- Python 3.8.12

### Additional context

<details>
<summary>Stacktrace</summary>

```
Traceback (most recent call last):
  File "/opt/project/src/scratch.py", line 30, in <module>
    trainer.validate(model, dataloader)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 848, in validate
    return self._call_and_handle_interrupt(self._validate_impl, model, dataloaders, ckpt_path, verbose, datamodule)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 721, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 895, in _validate_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1234, in _run
    results = self._run_stage()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1318, in _run_stage
    return self._run_evaluate()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1363, in _run_evaluate
    eval_loop_results = self._evaluation_loop.run()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 154, in advance
    dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 127, in advance
    output = self._evaluation_step(**kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 222, in _evaluation_step
    output = self.trainer._call_strategy_hook("validation_step", *kwargs.values())
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1763, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 344, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/double.py", line 52, in validation_step
    return self.module.validation_step(
  File "/opt/project/src/scratch.py", line 15, in validation_step
    print(self.metric(torch.tensor([50, 50, 50, 50]), torch.tensor([11, 9, 14, 22])))
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torchmetrics/metric.py", line 264, in forward
    self._forward_cache = self.compute()
  File "/opt/conda/lib/python3.8/site-packages/torchmetrics/metric.py", line 440, in wrapped_func
    value = compute(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torchmetrics/classification/precision_recall.py", line 303, in compute
    return _recall_compute(tp, fp, fn, self.average, self.mdmc_reduce)
  File "/opt/conda/lib/python3.8/site-packages/torchmetrics/functional/classification/precision_recall.py", line 258, in _recall_compute
    return _reduce_stat_scores(
  File "/opt/conda/lib/python3.8/site-packages/torchmetrics/functional/classification/stat_scores.py", line 264, in _reduce_stat_scores
    numerator = torch.where(zero_div_mask, tensor(float(zero_division), device=numerator.device), numerator)
RuntimeError: expected scalar type double but found float
```

</details>

Failing code snippet:
https://github.com/PyTorchLightning/metrics/blob/65ea713a7dcd0f9889fa34f54ba01c92c9baabd1/torchmetrics/functional/classification/stat_scores.py#L231-L264


Discussion:

Resolving patch:
diff --git a/torchmetrics/functional/classification/stat_scores.py b/torchmetrics/functional/classification/stat_scores.py
--- a/torchmetrics/functional/classification/stat_scores.py
+++ b/torchmetrics/functional/classification/stat_scores.py
@@ -261,9 +261,13 @@ def _reduce_stat_scores(
     else:
         weights = weights.float()
 
-    numerator = torch.where(zero_div_mask, tensor(float(zero_division), device=numerator.device), numerator)
-    denominator = torch.where(zero_div_mask | ignore_mask, tensor(1.0, device=denominator.device), denominator)
-    weights = torch.where(ignore_mask, tensor(0.0, device=weights.device), weights)
+    numerator = torch.where(
+        zero_div_mask, tensor(zero_division, dtype=numerator.dtype, device=numerator.device), numerator
+    )
+    denominator = torch.where(
+        zero_div_mask | ignore_mask, tensor(1.0, dtype=denominator.dtype, device=denominator.device), denominator
+    )
+    weights = torch.where(ignore_mask, tensor(0.0, dtype=weights.dtype, device=weights.device), weights)
 
     if average not in (AverageMethod.MICRO, AverageMethod.NONE, None):
         weights = weights / weights.sum(dim=-1, keepdim=True)
@@ -271,7 +275,7 @@ def _reduce_stat_scores(
     scores = weights * (numerator / denominator)
 
     # This is in case where sum(weights) = 0, which happens if we ignore the only present class with average='weighted'
-    scores = torch.where(torch.isnan(scores), tensor(float(zero_division), device=scores.device), scores)
+    scores = torch.where(torch.isnan(scores), tensor(zero_division, dtype=scores.dtype, device=scores.device), scores)
 
     if mdmc_average == MDMCAverageMethod.SAMPLEWISE:
         scores = scores.mean(dim=0)


Test patch:
diff --git a/tests/classification/test_stat_scores.py b/tests/classification/test_stat_scores.py
--- a/tests/classification/test_stat_scores.py
+++ b/tests/classification/test_stat_scores.py
@@ -173,10 +173,12 @@ class TestStatScores(MetricTester):
     # DDP tests temporarily disabled due to hanging issues
     @pytest.mark.parametrize("ddp", [False])
     @pytest.mark.parametrize("dist_sync_on_step", [True, False])
+    @pytest.mark.parametrize("dtype", [torch.float, torch.double])
     def test_stat_scores_class(
         self,
         ddp: bool,
         dist_sync_on_step: bool,
+        dtype: torch.dtype,
         sk_fn: Callable,
         preds: Tensor,
         target: Tensor,
@@ -191,6 +193,11 @@ def test_stat_scores_class(
         if ignore_index is not None and preds.ndim == 2:
             pytest.skip("Skipping ignore_index test with binary inputs.")
 
+        if preds.is_floating_point():
+            preds = preds.to(dtype)
+        if target.is_floating_point():
+            target = target.to(dtype)
+
         self.run_class_metric_test(
             ddp=ddp,
             preds=preds,