Repository: Lightning-AI/torchmetrics

Issue:
Can MultitaskWrapper.clone() have a postfix or prefix arg?
## 🚀 Feature
A `prefix` and `postfix` parameter to pass into the `clone()` method of `MultitaskWrapper`

### Motivation

I have fabricated a `MultitaskWrapper` to measure multiple task metrics for trainning, and wanted to clone it for validation and test.
```
self.metrics = MultitaskWrapper({
            'Img': MetricCollection({
                'PSNR': PeakSignalNoiseRatio(data_range=1.0),
                'SSIM': StructuralSimilarityIndexMeasure(data_range=1.0),
            }),
            'DoP': MetricCollection({`some_metrics`: SomeMetrics()}),
            'AoP': MetricCollection({`other_metrics`: OtherMetrics()})
        })
```
What I want to do is:
`self.val_metrics =  self.metrics.clone(postfix='_val')` so that we can obtain
```
self.val_metrics = MultitaskWrapper({
            'Img_val': MetricCollection({
                'PSNR': PeakSignalNoiseRatio(data_range=1.0),
                'SSIM': StructuralSimilarityIndexMeasure(data_range=1.0),
            }),
            'DoP_val': MetricCollection({`some_metrics`: SomeMetrics()}),
            'AoP_val': MetricCollection({`other_metrics`: OtherMetrics()})
        })
```
However, there is no such a postfix arg for `clone()`. So I can only re-define a `MultitaskWrapper` for `self.val_metrics` instead of simply cloning form `self.metrics`.

Is there any way to handle it?
Thx.



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

Resolving patch:
diff --git a/src/torchmetrics/wrappers/multitask.py b/src/torchmetrics/wrappers/multitask.py
--- a/src/torchmetrics/wrappers/multitask.py
+++ b/src/torchmetrics/wrappers/multitask.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # this is just a bypass for this module name collision with built-in one
+from copy import deepcopy
 from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 from torch import Tensor, nn
@@ -167,6 +168,33 @@ def reset(self) -> None:
             metric.reset()
         super().reset()
 
+    @staticmethod
+    def _check_arg(arg: Optional[str], name: str) -> Optional[str]:
+        if arg is None or isinstance(arg, str):
+            return arg
+        raise ValueError(f"Expected input `{name}` to be a string, but got {type(arg)}")
+
+    def clone(self, prefix: Optional[str] = None, postfix: Optional[str] = None) -> "MultitaskWrapper":
+        """Make a copy of the metric.
+
+        Args:
+            prefix: a string to append in front of the metric keys
+            postfix: a string to append after the keys of the output dict.
+
+        """
+        multitask_copy = deepcopy(self)
+        if prefix is not None:
+            prefix = self._check_arg(prefix, "prefix")
+            multitask_copy.task_metrics = nn.ModuleDict(
+                {prefix + key: value for key, value in multitask_copy.task_metrics.items()}
+            )
+        if postfix is not None:
+            postfix = self._check_arg(postfix, "postfix")
+            multitask_copy.task_metrics = nn.ModuleDict(
+                {key + postfix: value for key, value in multitask_copy.task_metrics.items()}
+            )
+        return multitask_copy
+
     def plot(
         self, val: Optional[Union[Dict, Sequence[Dict]]] = None, axes: Optional[Sequence[_AX_TYPE]] = None
     ) -> Sequence[_PLOT_OUT_TYPE]:


Test patch:
diff --git a/tests/unittests/wrappers/test_multitask.py b/tests/unittests/wrappers/test_multitask.py
--- a/tests/unittests/wrappers/test_multitask.py
+++ b/tests/unittests/wrappers/test_multitask.py
@@ -207,3 +207,18 @@ def test_nested_multitask_wrapper():
     multitask_results = multitask_metrics.compute()
 
     assert _dict_results_same_as_individual_results(classification_results, regression_results, multitask_results)
+
+
+def test_clone_with_prefix_and_postfix():
+    """Check that the clone method works with prefix and postfix arguments."""
+    multitask_metrics = MultitaskWrapper({"Classification": BinaryAccuracy(), "Regression": MeanSquaredError()})
+    cloned_metrics_with_prefix = multitask_metrics.clone(prefix="prefix_")
+    cloned_metrics_with_postfix = multitask_metrics.clone(postfix="_postfix")
+
+    # Check if the cloned metrics have the expected keys
+    assert set(cloned_metrics_with_prefix.task_metrics.keys()) == {"prefix_Classification", "prefix_Regression"}
+    assert set(cloned_metrics_with_postfix.task_metrics.keys()) == {"Classification_postfix", "Regression_postfix"}
+
+    # Check if the cloned metrics have the expected values
+    assert isinstance(cloned_metrics_with_prefix.task_metrics["prefix_Classification"], BinaryAccuracy)
+    assert isinstance(cloned_metrics_with_prefix.task_metrics["prefix_Regression"], MeanSquaredError)