Repository: Lightning-AI/torchmetrics

Issue:
`compute_on_cpu=True` not respected
## 🐛 Bug

I am using `torchmetrics.classification.roc.ROC` as a parent class in another project [1], and I had OOM issues because of the accumulation of preds and targes (when `threshold=None`). 

I tried using `compute_on_cpu=True` but the class doesn't seem to respect the option.

**Am I missing something?**
`preds.device` at training_step: cuda:0
self.auroc.preds[-1].device: cuda:0
self.auroc.compute().device: cuda:0
[1] The project is [`anomalib`](https://github.com/openvinotoolkit/anomalib), and the class inheriting from `ROC` is [`anomalib.utils.metrics.auroc.AUROC`](https://github.com/openvinotoolkit/anomalib/blob/74595c9a5a519cdc20fad1f3c3c68b5d0434cf44/anomalib/utils/metrics/auroc.py#L17) (yes, it shadows 'torchmetrics.AUROC').

---

I took screenshots of what's happenning inside the class `ROC` with vscode debugger. The training is done with pytorchlightning's trainer.

### `__init__()`

`self.compute_on_cpu` is correctly `True`.

![debug init](https://user-images.githubusercontent.com/24547377/200527724-36291483-68fc-4d08-a13d-4f7b555db55e.png)

### `update()`

Both `preds` (func arg) and `self.preds` (class attribute) are on the cpu, as expected.

Same for `target`/`self.target`.

![debug update](https://user-images.githubusercontent.com/24547377/200527941-09bd817b-5f75-48c4-9944-10a6444dc25d.png)

### `compute()`

Both `preds` (func arg) and `self.preds` (class attribute) **were sent to the GPU** -- which I did **NOT** want.

Same for `target`/`self.target`.

![debug compute](https://user-images.githubusercontent.com/24547377/200528130-304d6998-1690-44b6-8794-11dadde92318.png)

### To Reproduce

I got the [`mnist-hello-world.ipynb`](https://colab.research.google.com/github/PytorchLightning/lightning-tutorials/blob/publication/.notebooks/lightning_examples/mnist-hello-world.ipynb) from pytorchlightning's tutorials and modified it a bit to use that class in my project.

The notebook can be seen here: [github.com/jpcbertoldo/92eb5c91c4e8664bcc91ea34a4d9d37a](https://gist.github.com/jpcbertoldo/92eb5c91c4e8664bcc91ea34a4d9d37a).

The prints inside the `training_step()` show this (`FORCE_PREDS_TO_CPU = False`):

```text
`preds.device` at training_step: cuda:0
self.auroc.preds[-1].device: cuda:0
self.auroc.compute().device: cuda:0
```

I.e., the `compute_on_cpu` is not being respected -- right?

#### Code sample

Here is a simplified copy of `AUROC` from `anomalib`. I believe it's **not** influencing much because it's very simple and one doesn't see any `.to(device)`.

```py
class AUROC(ROC):
    """Area under the ROC curve."""

    def compute(self) -> Tensor:
        fpr, tpr = self._compute()
        # TODO: use stable sort after upgrading to pytorch 1.9.x (https://github.com/openvinotoolkit/anomalib/issues/92)
        if not (torch.all(fpr.diff() <= 0) or torch.all(fpr.diff() >= 0)):
            return auc(fpr, tpr, reorder=True)  # only reorder if fpr is not increasing or decreasing
        return auc(fpr, tpr)

    def update(self, preds: Tensor, target: Tensor) -> None:
        super().update(preds.flatten(), target.flatten())

    def _compute(self) -> Tuple[Tensor, Tensor]:
        fpr, tpr, _thresholds = super().compute()
        return (fpr, tpr)
```

### Expected behavior

I want the metric to be computed on the CPU and whatever internal tensor should stay in the CPU all times.

### Environment

Environment installed with miniconda, but the relevant dependencies are installed via pip (in the conda env file).

```
Python 3.8.13
pytorch-lightning             1.6.5
torch                         1.12.0
torchmetrics                  0.9.3
torchvision                   0.13.0
```



Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -288,6 +288,8 @@ def _forward_full_state_update(self, *args: Any, **kwargs: Any) -> Any:
         self._computed = None
         self._enable_grad = False
         self.compute_on_cpu = _temp_compute_on_cpu
+        if self.compute_on_cpu:
+            self._move_list_states_to_cpu()
 
         return batch_val
 
@@ -325,6 +327,8 @@ def _forward_reduce_state_update(self, *args: Any, **kwargs: Any) -> Any:
         self._computed = None
         self._enable_grad = False
         self.compute_on_cpu = _temp_compute_on_cpu
+        if self.compute_on_cpu:
+            self._move_list_states_to_cpu()
 
         return batch_val
 


Test patch:
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
@@ -471,3 +471,19 @@ def test_no_iteration_allowed():
     with pytest.raises(NotImplementedError, match="Metrics does not support iteration."):
         for m in metric:
             continue
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires cuda")
+@pytest.mark.parametrize("method", ["forward", "update"])
+def test_compute_on_cpu_arg_forward(method):
+    metric = DummyListMetric(compute_on_cpu=True)
+    x = torch.randn(10).cuda()
+    if method == "update":
+        metric.update(x)
+        metric.update(x)
+    else:
+        _ = metric(x)
+        _ = metric(x)
+    val = metric.compute()
+    assert all(str(v.device) == "cpu" for v in val)
+    assert all(torch.allclose(v, x.cpu()) for v in val)