Repository: Lightning-AI/torchmetrics

Issue:
Unable to plot `MetricCollection` containing prefix using `MetricCollection.plot()`
## 🐛 Bug

I want to evaluate and plot the same metrics, namely confusion matrix and ROC curve, during training, validation and testing phase. For this, I set up a base `MetricCollection` which is then cloned three times. The respective prefix, e.g. `'train_'`, is added using the prefix argument of the `MetricCollection.clone()` function.

When trying to plot the `MetricCollection` with the prefix using `MetricCollection.plot()`, a `KeyError` exception is thrown. When no prefix is provided, `MetricCollection.plot()` works correctly.

### To Reproduce

Steps to reproduce the behavior...

**Fails: with prefix**
<details>
  <summary>Code sample</summary>

```python
import torchmetrics
import torch
import matplotlib.pyplot as plt

#test data
preds = torch.tensor([[ 0.,  0.,  1.],
                      [ 1.,  0.,  0.],
                      [ 0.,  1.,  0.],
                      [ 0.,  0.,  1.],
                      [ 0.,  0.,  1.],
                      [ 1.,  0.,  0.]])
target = torch.tensor([2, 0, 1, 2, 0, 1])
# base metric collection
conf_and_roc = torchmetrics.MetricCollection([torchmetrics.ROC(task="multiclass", num_classes=3),
                                              torchmetrics.ConfusionMatrix(task="multiclass", num_classes=3)])
# clone base metric collection and add prefix
train_conf_and_roc = conf_and_roc.clone(prefix='train_')
# update metrics in collection
train_conf_and_roc.update(preds, target)
# plot metrics
res = train_conf_and_roc.plot()
# show resulting figures
plt.show()
```
</details>

<details>
  <summary>Error message</summary>

```bash
File "/Users/daniel/lightning_minimal_example/demo.py", line 21, in <module>
    res = train_conf_and_roc.plot()
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/daniel/anaconda3/envs/lightning_example/lib/python3.11/site-packages/torchmetrics/collections.py", line 657, in plot
    f, a = m.plot(val[k], ax=ax[i] if ax is not None else ax)
                  ~~~^^^
KeyError: 'MulticlassROC'
```
</details>



**Works: without prefix**
<details>
  <summary>Code sample</summary>

```python
import torchmetrics
import torch
import matplotlib.pyplot as plt

#test data
preds = torch.tensor([[ 0.,  0.,  1.],
                      [ 1.,  0.,  0.],
                      [ 0.,  1.,  0.],
                      [ 0.,  0.,  1.],
                      [ 0.,  0.,  1.],
                      [ 1.,  0.,  0.]])
target = torch.tensor([2, 0, 1, 2, 0, 1])
# base metric collection
conf_and_roc = torchmetrics.MetricCollection([torchmetrics.ROC(task="multiclass", num_classes=3),
                                              torchmetrics.ConfusionMatrix(task="multiclass", num_classes=3)])
# clone base metric collection and add prefix
train_conf_and_roc = conf_and_roc.clone()
# update metrics in collection
train_conf_and_roc.update(preds, target)
# plot metrics
res = train_conf_and_roc.plot()
# show resulting figures
plt.show()
```
</details>

### Expected behavior

No `KeyError` when plotting a `MetricCollection` with a set prefix property.

### Environment

- OS: macOS Sonoma 14.3.1 
- python: 3.11.5
- torch: 2.2.1
- torchmetrics: 1.3.1


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

Resolving patch:
diff --git a/src/torchmetrics/collections.py b/src/torchmetrics/collections.py
--- a/src/torchmetrics/collections.py
+++ b/src/torchmetrics/collections.py
@@ -647,12 +647,11 @@ def plot(
                     f"Expected argument `ax` to be a sequence of matplotlib axis objects with the same length as the "
                     f"number of metrics in the collection, but got {type(ax)} with len {len(ax)} when `together=False`"
                 )
-
         val = val or self.compute()
         if together:
             return plot_single_or_multi_val(val, ax=ax)
         fig_axs = []
-        for i, (k, m) in enumerate(self.items(keep_base=True, copy_state=False)):
+        for i, (k, m) in enumerate(self.items(keep_base=False, copy_state=False)):
             if isinstance(val, dict):
                 f, a = m.plot(val[k], ax=ax[i] if ax is not None else ax)
             elif isinstance(val, Sequence):


Test patch:
diff --git a/tests/unittests/utilities/test_plot.py b/tests/unittests/utilities/test_plot.py
--- a/tests/unittests/utilities/test_plot.py
+++ b/tests/unittests/utilities/test_plot.py
@@ -834,12 +834,17 @@ def test_confusion_matrix_plotter(metric_class, preds, target, labels, use_label
 
 @pytest.mark.parametrize("together", [True, False])
 @pytest.mark.parametrize("num_vals", [1, 2])
-def test_plot_method_collection(together, num_vals):
+@pytest.mark.parametrize(
+    ("prefix", "postfix"), [(None, None), ("prefix", None), (None, "postfix"), ("prefix", "postfix")]
+)
+def test_plot_method_collection(together, num_vals, prefix, postfix):
     """Test the plot method of metric collection."""
     m_collection = MetricCollection(
         BinaryAccuracy(),
         BinaryPrecision(),
         BinaryRecall(),
+        prefix=prefix,
+        postfix=postfix,
     )
     if num_vals == 1:
         m_collection.update(torch.randint(0, 2, size=(10,)), torch.randint(0, 2, size=(10,)))