Repository: Lightning-AI/torchmetrics

Issue:
DeepSpeed still changes metric states from fp32 to fp16
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
Following #484, the PR #493 introduced `set_dtype()` to prevent `.half()` calls to change the precision of metric states. However, at least for `PearsonCorrCoef`, DeepSpeed still somehow modifies the `dtype`. During initialization the metric states have the default dtype of `torch.float32`. However, this changes as soon as DeepSpeed is initialized (refer to the code example below).

I tried this with `MeanAbsoluteError` as well, and interestingly, out of `sum_abs_error` and `total` metric states, only `sum_abs_error` changes from `torch.float32` to `torch.float16`. The `total` metric state remains as `torch.int64`. So DeepSpeed is probably only converting floats and not ints. 

This is especially problematic for epoch level metrics such as `PearsonCorrCoef`, since the numbers they hold can easily become larger the the 65,536 maximum allowed with fp16 precision. 
### To Reproduce

Refer to the minimal code sample.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

<details>
  <summary>Code sample</summary>

```python
import os
import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer, LightningDataModule
from torchmetrics import PearsonCorrCoef, MeanAbsoluteError


class RandomDataset(Dataset):
    def __init__(self, size, num_samples):
        self.len = num_samples
        self.data = torch.randn(num_samples, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class PlDataModule(LightningDataModule):
    def train_dataloader(self):
        return DataLoader(RandomDataset(32, 64), batch_size=2)


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 32)
        self.metric = PearsonCorrCoef()
        self.mae = MeanAbsoluteError()
        print("Before DeepSpeed initialization")
        print("self.metric.mean_x", self.metric.mean_x)
        print("self.metric.mean_x.dtype", self.metric.mean_x.dtype)
        print("self.metric.mean_y", self.metric.mean_y)
        print("self.metric.dtype", self.metric.mean_y.dtype)
        print("self.metric.var_x", self.metric.var_x)
        print("self.metric.var_x.dtype", self.metric.var_x.dtype)
        print("self.metric.var_y", self.metric.var_y)
        print("self.metric.var_y.dtype", self.metric.var_y.dtype)
        print("self.metric.corr_xy", self.metric.corr_xy)
        print("self.metric.corr_xy.dtype", self.metric.corr_xy.dtype)
        print("self.metric.n_total", self.metric.n_total)

        print("self.mae.sum_abs_error.dtype", self.mae.sum_abs_error.dtype)
        print("self.mae.total.dtype", self.mae.total.dtype)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        pred = self.forward(batch)
        loss = self(batch).sum()
        self.metric.update(torch.flatten(pred), torch.flatten(batch))

        print("After DeepSpeed initialization")
        print("self.metric.mean_x", self.metric.mean_x)
        print("self.metric.mean_x.dtype", self.metric.mean_x.dtype)
        print("self.metric.mean_y", self.metric.mean_y)
        print("self.metric.dtype", self.metric.mean_y.dtype)
        print("self.metric.var_x", self.metric.var_x)
        print("self.metric.var_x.dtype", self.metric.var_x.dtype)
        print("self.metric.var_y", self.metric.var_y)
        print("self.metric.var_y.dtype", self.metric.var_y.dtype)
        print("self.metric.corr_xy", self.metric.corr_xy)
        print("self.metric.corr_xy.dtype", self.metric.corr_xy.dtype)
        print("self.metric.n_total", self.metric.n_total)

        print("self.mae.sum_abs_error.dtype", self.mae.sum_abs_error.dtype)
        print("self.mae.total.dtype", self.mae.total.dtype)

        return {"loss": loss}

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    model = BoringModel()
    trainer = Trainer(
        default_root_dir="/scratch/l/lstein/ftaj/scTransformers/scTransPert/checkpoints/temp/",
        # default_root_dir=os.getcwd(),
        limit_train_batches=10,
        num_sanity_val_steps=0,
        max_epochs=1,
        strategy="deepspeed_stage_1",
        accelerator="gpu",
        precision=16
    )
    trainer.fit(model, datamodule=PlDataModule())


run()
```

</details>

<details><summary>Output on my setup</summary>

```
(autopath) [ftaj@mist021 autopath]$ python src/deepspeed_issue.py
Before DeepSpeed initialization
self.metric.mean_x tensor([0.])
self.metric.mean_x.dtype torch.float32
self.metric.mean_y tensor([0.])
self.metric.dtype torch.float32
self.metric.var_x tensor([0.])
self.metric.var_x.dtype torch.float32
self.metric.var_y tensor([0.])
self.metric.var_y.dtype torch.float32
self.metric.corr_xy tensor([0.])
self.metric.corr_xy.dtype torch.float32
self.metric.n_total tensor([0.])
self.mae.sum_abs_error.dtype torch.float32
self.mae.total.dtype torch.int64
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
initializing deepspeed distributed: GLOBAL_RANK: 0, MEMBER: 1/1
Enabling DeepSpeed FP16.
/home/l/lstein/ftaj/.conda/envs/autopath/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:613: UserWarning: Checkpoint directory /scratch/l/lstein/ftaj/scTransformers/scTransPert/checkpoints/temp/lightning_logs/version_357072/checkpoints exists and is not empty.
  rank_zero_warn(f"Checkpoint directory {dirpath} exists and is not empty.")
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
[2023-02-26 21:48:53,595] [WARNING] [engine.py:1223:_do_optimizer_sanity_check] **** You are using ZeRO with an untested optimizer, proceed with caution *****
Rank: 0 partition count [1] and sizes[(1056, False)]

  | Name   | Type              | Params
---------------------------------------------
0 | layer  | Linear            | 1.1 K
1 | metric | PearsonCorrCoef   | 0
2 | mae    | MeanAbsoluteError | 0
---------------------------------------------
1.1 K     Trainable params
0         Non-trainable params
1.1 K     Total params
0.002     Total estimated model params size (MB)
SLURM auto-requeueing enabled. Setting signal handlers.
/home/l/lstein/ftaj/.conda/envs/autopath/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 128 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
/home/l/lstein/ftaj/.conda/envs/autopath/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1609: PossibleUserWarning: The number of training batches (10) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
  rank_zero_warn(
Epoch 0:   0%|                                                                                                                                                                                                                                           | 0/10 [00:00<?, ?it/s]
After DeepSpeed initialization
self.metric.mean_x tensor([-0.0167], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0883], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([18.6094], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([71.9375], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([4.5664], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([64.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  10%|████████████████████                                                                                                                                                                                     | 1/10 [00:00<00:04,  2.20it/s, loss=-1.07, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([0.0207], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0510], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([34.7812], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([135.1250], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([-1.7344], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([128.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  20%|████████████████████████████████████████▍                                                                                                                                                                 | 2/10 [00:00<00:01,  4.31it/s, loss=1.33, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([0.0435], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0620], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([55.9375], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([202.5000], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([-3.5781], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([192.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  30%|████████████████████████████████████████████████████████████▌                                                                                                                                             | 3/10 [00:00<00:01,  6.35it/s, loss=2.78, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([0.0287], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0602], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([73.5625], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([258.5000], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([-11.3281], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([256.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  40%|████████████████████████████████████████████████████████████████████████████████▊                                                                                                                         | 4/10 [00:00<00:00,  8.29it/s, loss=1.83, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([-0.0767], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0827], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([107.2500], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([315.2500], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([-8.4688], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([320.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  50%|████████████████████████████████████████████████████████████████████████████████████████████████████▌                                                                                                    | 5/10 [00:00<00:00, 10.15it/s, loss=-4.91, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([-0.0933], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0465], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([132.8750], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([394.5000], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([-4.7578], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([384.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  60%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                                                                                | 6/10 [00:00<00:00, 11.97it/s, loss=-5.97, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([-0.3110], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0288], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([371.5000], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([456.2500], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([-15.1953], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([448.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  70%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋                                                            | 7/10 [00:00<00:00, 13.70it/s, loss=-19.9, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([-0.1224], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0084], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([673.], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([510.5000], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([8.4297], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([512.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  80%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊                                        | 8/10 [00:00<00:00, 15.37it/s, loss=-7.84, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([-0.1087], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([-0.0055], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([739.], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([568.], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([1.3281], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([576.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0:  90%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉                    | 9/10 [00:00<00:00, 16.98it/s, loss=-6.96, v_num=357072]
After DeepSpeed initialization
self.metric.mean_x tensor([-0.0939], device='cuda:0', dtype=torch.float16)
self.metric.mean_x.dtype torch.float16
self.metric.mean_y tensor([0.0134], device='cuda:0', dtype=torch.float16)
self.metric.dtype torch.float16
self.metric.var_x tensor([862.5000], device='cuda:0', dtype=torch.float16)
self.metric.var_x.dtype torch.float16
self.metric.var_y tensor([642.], device='cuda:0', dtype=torch.float16)
self.metric.var_y.dtype torch.float16
self.metric.corr_xy tensor([-19.6406], device='cuda:0', dtype=torch.float16)
self.metric.corr_xy.dtype torch.float16
self.metric.n_total tensor([640.], device='cuda:0', dtype=torch.float16)
self.mae.sum_abs_error.dtype torch.float16
self.mae.total.dtype torch.int64
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 18.49it/s, loss=-6.01, v_num=357072]/home/l/lstein/ftaj/.conda/envs/autopath/lib/python3.8/site-packages/torch/nn/modules/module.py:1428: UserWarning: Positional args are being deprecated, use kwargs instead. Refer to https://pytorch.org/docs/master/generated/torch.nn.Module.html#torch.nn.Module.state_dict for details.
  warnings.warn(
`Trainer.fit` stopped: `max_epochs=1` reached.
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 17.59it/s, loss=-6.01, v_num=357072]
```

</details>


<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->



### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
DeepSpeed should not affect `dtype` of metric states, that is, they should remain in `torch.float32` even if DeepSpeed with `precision=16` is used.
### Environment
<details><summary>collect_env results</summary>

```
* CUDA:
        - GPU:
                - Tesla V100-SXM2-32GB
        - available:         True
        - version:           11.4
* Lightning:
        - lightning-cloud:   0.5.26
        - lightning-utilities: 0.6.0.post0
        - pytorch-lightning: 1.9.1
        - reformer-pytorch:  1.4.4
        - torch:             1.13.0

        - av:                10.0.0
        - axial-positional-embedding: 0.2.1
        - backcall:          0.2.0
        - beautifulsoup4:    4.11.1
        - bleach:            4.1.0
        - blessed:           1.20.0
        - bottleneck:        1.3.5
        - brotlipy:          0.7.0
        - cached-property:   1.5.2
        - captum:            0.5.0
        - certifi:           2022.12.7
        - cffi:              1.15.1
        - charset-normalizer: 2.0.4
        - click:             8.0.4
        - cmake:             3.25.0
        - croniter:          1.3.8
        - cryptography:      38.0.1
        - cycler:            0.11.0
        - dateutils:         0.6.12
        - debugpy:           1.5.1
        - decorator:         5.1.1
        - deepdiff:          6.2.3
        - deepspeed:         0.8.0
        - defusedxml:        0.7.1
        - dnspython:         2.3.0
        - docker-pycreds:    0.4.0
        - einops:            0.4.1
        - email-validator:   1.3.1
        - entrypoints:       0.4
        - executing:         0.8.3
        - fastapi:           0.88.0
        - fastjsonschema:    2.16.2
        - filelock:          3.8.2
        - fire:              0.4.0
        - frozenlist:        1.3.3
        - fsspec:            2022.11.0
        - future:            0.18.2
        - gitdb:             4.0.9
        - gitpython:         3.1.29
        - h11:               0.14.0
        - h5py:              3.6.0
        - hjson:             3.1.0
        - httpcore:          0.16.3
        - httptools:         0.5.0
        - httpx:             0.23.3
        - idna:              3.4
        - importlib-metadata: 5.0.0
        - importlib-resources: 5.2.0
        - inquirer:          3.1.2
        - ipykernel:         6.15.2
        - ipython:           8.6.0
        - ipython-genutils:  0.2.0
        - itsdangerous:      2.1.2
        - jedi:              0.18.1
        - jinja2:            3.1.2
        - joblib:            1.2.0
        - jsonschema:        4.16.0
        - jupyter-client:    7.3.4
        - jupyter-core:      4.11.2

        - readchar:          4.0.3
        - reformer-pytorch:  1.4.4
        - requests:          2.28.1
        - rfc3986:           1.5.0
        - rich:              13.3.1
        - routing-transformer: 1.6.1
        - scanpy:            1.9.1
        - scikit-learn:      1.1.2
        - scipy:             1.8.1
        - seaborn:           0.12.1
        - send2trash:        1.8.0
        - sentencepiece:     0.1.96
        - sentry-sdk:        1.10.1
        - session-info:      1.0.0
        - setproctitle:      1.3.2
        - setuptools:        65.5.0
        - shortuuid:         1.0.10
        - six:               1.16.0
        - smmap:             5.0.0
        - sniffio:           1.2.0
        - soupsieve:         2.3.2.post1
        - stack-data:        0.2.0
        - starlette:         0.22.0
        - starsessions:      1.3.0
        - statsmodels:       0.11.1
        - stdlib-list:       0.8.0
        - tabulate:          0.8.10
        - tensorboardx:      2.5.1
        - termcolor:         2.1.1
        - terminado:         0.13.1
        - threadpoolctl:     3.1.0
        - tinycss2:          1.2.1
        - torch:             1.13.0
        - torchmetrics:      0.11.2
        - torchtext:         0.14.0a0+f8d65dd
        - torchvision:       0.14.0
        - tornado:           6.1
        - tqdm:              4.64.1
        - traitlets:         5.9.0
        - typing-extensions: 4.3.0
        - typing-inspect:    0.8.0
        - ujson:             5.7.0
        - umap-learn:        0.5.2
        - urllib3:           1.26.12
        - uvicorn:           0.20.0
        - uvloop:            0.17.0
        - wandb:             0.13.5
        - watchfiles:        0.18.1
        - wcwidth:           0.2.5
        - webencodings:      0.5.1
        - websocket-client:  0.58.0
        - websockets:        10.4
        - wheel:             0.37.1
        - x-transformers:    1.5.1
        - xformers:          0.0.13
        - yarl:              1.8.2
        - zipp:              3.10.0
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - ELF
        - processor:         ppc64le
        - python:            3.8.13
        - version:           #1 SMP Thu Nov 17 08:57:02 EST 2022
```

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
@justusschock do you have experience with the internals of DeepSpeed? Is this something we can get around or is this a limitation?
@SkafteNicki I don't have too much experience with deepspeed internals, but I'll have a look why/where this happens.
@justusschock Hi! One thing I noticed is that the `PearsonCorrCoef` shows up in the model summary, even though I'm not using this metric for backpropagation, only during validation. Since `PearsonCorrCoef` is differentiable, perhaps it gets added to the computation graph and DeepSpeed finds another way to convert it? Just a guess.
Hey @FarzanT , `PersonCorrCoef` shows up in the model summary because it is a subclass of `Metric` (which in turn is a subclass of `torch.nn.Module`) and model summary lists all nn.Modules attached to the lightning module. This should, however, not result in devices changes for metrics as this is explicitly guarded.
@justusschock 
I realized calling `.half()` on a lightning module converts the metric states. This can be seen with the following code example:

```python
import torch

from pytorch_lightning import LightningModule
from torchmetrics import PearsonCorrCoef, MeanAbsoluteError


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 32)
        self.pearson = PearsonCorrCoef()
        self.mae = MeanAbsoluteError()

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        pred = self.forward(batch)
        loss = self(batch).sum()
        self.metric.update(torch.flatten(pred), torch.flatten(batch))

        return {"loss": loss}

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)

model = BoringModel()
print(model.mae.sum_abs_error.dtype)
print(model.pearson.mean_x.dtype)
model = model.half()
print(model.mae.sum_abs_error.dtype)
print(model.pearson.mean_x.dtype)
model = model.float()
print(model.mae.sum_abs_error.dtype)
print(model.pearson.mean_x.dtype)
```
Output:
```
torch.float32
torch.float32
torch.float16
torch.float16
torch.float32
torch.float32
```

Any pointers on how to protect the metrics states from this?
@justusschock @SkafteNicki  I think I figured out the source of the problem:
https://github.com/Lightning-AI/metrics/blob/bc722a7e487628a8be83bd8e4eaddc25c18698b2/src/torchmetrics/metric.py#L697-L730
The `_apply()` function in `metric.py` was specifically made to apply `.to()`, `.cuda()` etc calls on the metric states. I guess the same goes for `.half()`. So, when you call `.half()` on a Lightning Module, `.half()` is applied directly on the metric states, circumventing the 'guards' proposed in #493. 
Should we put a check in the `_apply()` function to ignore functions that change dtype, like `half()` and `float()` etc.?

May I open a pull request?

Resolving patch:
diff --git a/src/torchmetrics/collections.py b/src/torchmetrics/collections.py
--- a/src/torchmetrics/collections.py
+++ b/src/torchmetrics/collections.py
@@ -472,3 +472,13 @@ def __repr__(self) -> str:
         if self.postfix:
             repr_str += f"{',' if not self.prefix else ''}\n  postfix={self.postfix}"
         return repr_str + "\n)"
+
+    def set_dtype(self, dst_type: Union[str, torch.dtype]) -> "MetricCollection":
+        """Transfer all metric state to specific dtype. Special version of standard `type` method.
+
+        Arguments:
+            dst_type (type or string): the desired type.
+        """
+        for _, m in self.items(keep_base=True, copy_state=False):
+            m.set_dtype(dst_type)
+        return self
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -145,6 +145,7 @@ def __init__(
         self._to_sync = self.sync_on_compute
         self._should_unsync = True
         self._enable_grad = False
+        self._dtype_convert = False
 
         # initialize state
         self._defaults: Dict[str, Union[List, Tensor]] = {}
@@ -692,14 +693,23 @@ def set_dtype(self, dst_type: Union[str, torch.dtype]) -> "Metric":
         Arguments:
             dst_type (type or string): the desired type.
         """
-        return super().type(dst_type)
+        self._dtype_convert = True
+        out = super().type(dst_type)
+        out._dtype_convert = False
+        return out
 
     def _apply(self, fn: Callable) -> Module:
         """Overwrite _apply function such that we can also move metric states to the correct device.
 
-        This method is called by the base ``nn.Module`` class whenever `.to`, `.cuda`, etc. methods are called.
+        This method is called by the base ``nn.Module`` class whenever `.to`, `.cuda`, `.float`, `.half` etc. methods
+        are called. Dtype conversion is garded and will only happen through the special `set_dtype` method.
         """
         this = super()._apply(fn)
+        fs = str(fn)
+        cond = any(f in fs for f in ["Module.type", "Module.half", "Module.float", "Module.double", "Module.bfloat16"])
+        if not self._dtype_convert and cond:
+            return this
+
         # Also apply fn to metric states and defaults
         for key, value in this._defaults.items():
             if isinstance(value, Tensor):


Test patch:
diff --git a/tests/integrations/test_lightning.py b/tests/integrations/test_lightning.py
--- a/tests/integrations/test_lightning.py
+++ b/tests/integrations/test_lightning.py
@@ -16,6 +16,7 @@
 import torch
 from pytorch_lightning import LightningModule, Trainer
 from torch import tensor
+from torch.nn import Linear
 from torch.utils.data import DataLoader
 
 from integrations.helpers import no_warning_call
@@ -305,3 +306,48 @@ def training_step(self, batch, batch_idx):
     output = model(rand_input)
     script_output = script_model(rand_input)
     assert torch.allclose(output, script_output)
+
+
+def test_dtype_in_pl_module_transfer(tmpdir):
+    """Test that metric states don't change dtype when .half() or .float() is called on the LightningModule."""
+
+    class BoringModel(LightningModule):
+        def __init__(self, metric_dtype=torch.float32):
+            super().__init__()
+            self.layer = Linear(32, 32)
+            self.metric = SumMetric()
+            self.metric.set_dtype(metric_dtype)
+
+        def forward(self, x):
+            return self.layer(x)
+
+        def training_step(self, batch, batch_idx):
+            pred = self.forward(batch)
+            loss = self(batch).sum()
+            self.metric.update(torch.flatten(pred), torch.flatten(batch))
+
+            return {"loss": loss}
+
+        def configure_optimizers(self):
+            return torch.optim.SGD(self.layer.parameters(), lr=0.1)
+
+    model = BoringModel()
+    assert model.metric.value.dtype == torch.float32
+    model = model.half()
+    assert model.metric.value.dtype == torch.float32
+
+    model = BoringModel()
+    assert model.metric.value.dtype == torch.float32
+    model = model.double()
+    assert model.metric.value.dtype == torch.float32
+
+    model = BoringModel(metric_dtype=torch.float16)
+    assert model.metric.value.dtype == torch.float16
+    model = model.float()
+    assert model.metric.value.dtype == torch.float16
+
+    model = BoringModel()
+    assert model.metric.value.dtype == torch.float32
+
+    model = model.type(torch.half)
+    assert model.metric.value.dtype == torch.float32
diff --git a/tests/unittests/bases/test_collections.py b/tests/unittests/bases/test_collections.py
--- a/tests/unittests/bases/test_collections.py
+++ b/tests/unittests/bases/test_collections.py
@@ -93,11 +93,11 @@ def test_device_and_dtype_transfer_metriccollection(tmpdir):
     for _, metric in metric_collection.items():
         assert metric.x.is_cuda
 
-    metric_collection = metric_collection.double()
+    metric_collection = metric_collection.set_dtype(torch.double)
     for _, metric in metric_collection.items():
         assert metric.x.dtype == torch.float64
 
-    metric_collection = metric_collection.half()
+    metric_collection = metric_collection.set_dtype(torch.half)
     for _, metric in metric_collection.items():
         assert metric.x.dtype == torch.float16
 
diff --git a/tests/unittests/bases/test_metric.py b/tests/unittests/bases/test_metric.py
--- a/tests/unittests/bases/test_metric.py
+++ b/tests/unittests/bases/test_metric.py
@@ -310,6 +310,21 @@ def test_device_and_dtype_transfer(tmpdir):
     assert metric.x.dtype == torch.float16
 
 
+def test_disable_of_normal_dtype_methods():
+    """Check that the default dtype changing methods does nothing."""
+    metric = DummyMetricSum()
+    assert metric.x.dtype == torch.float32
+
+    metric = metric.half()
+    assert metric.x.dtype == torch.float32
+
+    metric = metric.double()
+    assert metric.x.dtype == torch.float32
+
+    metric = metric.type(torch.half)
+    assert metric.x.dtype == torch.float32
+
+
 def test_warning_on_compute_before_update():
     """Test that an warning is raised if user tries to call compute before update."""
     metric = DummyMetricSum()