---
title: Multi-task Wrapper
description: >-
  Wrapper class for computing different metrics on different tasks in the
  context of multitask learning.
canonical: 'https://lightning.ai/docs/torchmetrics/stable/wrappers/multi_task_wrapper'
---
# Multi-task Wrapper {#multi-task-wrapper}

## Module Interface {#module-interface}

### `class torchmetrics.wrappers.MultitaskWrapper(task_metrics, prefix=None, postfix=None)` {#torchmetrics.wrappers.MultitaskWrapper}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L31-L389)

Wrapper class for computing different metrics on different tasks in the context of multitask learning.

In multitask learning the different tasks requires different metrics to be evaluated. This wrapper allows for easy evaluation in such cases by supporting multiple predictions and targets through a dictionary. Note that only metrics where the signature of update follows the standard preds, target is supported.

Parameters:

*   **task\_metrics** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), `Union`\[[`Metric`](/docs/torchmetrics/1.9.0/api-reference/metric.md#torchmetrics.Metric "torchmetrics.metric.Metric"), [`MetricCollection`](/docs/torchmetrics/1.9.0/user-guide/overview.md#torchmetrics.MetricCollection "torchmetrics.collections.MetricCollection")\]\]) – Dictionary associating each task to a Metric or a MetricCollection. The keys of the dictionary represent the names of the tasks, and the values represent the metrics to use for each task.
    
*   **prefix** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]) – A string to append in front of the metric keys. If not provided, will default to an empty string.
    
*   **postfix** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]) – A string to append after the keys of the output dict. If not provided, will default to an empty string.
    

Tip

The use prefix and postfix allows for easily creating task wrappers for training, validation and test. The arguments are only changing the output keys of the computed metrics and not the input keys. This means that a `MultitaskWrapper` initialized as `MultitaskWrapper({"task": Metric()}, prefix="train_")` will still expect the input to be a dictionary with the key “task”, but the output will be a dictionary with the key “train\_task”.

Raises:

*   [**TypeError**](https://docs.python.org/3/builtins/exceptions.html#TypeError "(in Python v3.15)") – If argument task\_metrics is not an dictionary
    
*   [**TypeError**](https://docs.python.org/3/builtins/exceptions.html#TypeError "(in Python v3.15)") – If not all values in the task\_metrics dictionary is instances of Metric or MetricCollection
    
*   [**ValueError**](https://docs.python.org/3/builtins/exceptions.html#ValueError "(in Python v3.15)") – If prefix is not a string
    
*   [**ValueError**](https://docs.python.org/3/builtins/exceptions.html#ValueError "(in Python v3.15)") – If postfix is not a string
    

Example (with a single metric per class):

```python
>>> import torch
>>> from torchmetrics.wrappers import MultitaskWrapper
>>> from torchmetrics.regression import MeanSquaredError
>>> from torchmetrics.classification import BinaryAccuracy
>>>
>>> classification_target = torch.tensor([0, 1, 0])
>>> regression_target = torch.tensor([2.5, 5.0, 4.0])
>>> targets = {"Classification": classification_target, "Regression": regression_target}
>>>
>>> classification_preds = torch.tensor([0, 0, 1])
>>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
>>> preds = {"Classification": classification_preds, "Regression": regression_preds}
>>>
>>> metrics = MultitaskWrapper({
...     "Classification": BinaryAccuracy(),
...     "Regression": MeanSquaredError()
... })
>>> metrics.update(preds, targets)
>>> metrics.compute()
{'Classification': tensor(0.3333), 'Regression': tensor(0.8333)}
```

Example (with several metrics per task):

```python
>>> import torch
>>> from torchmetrics import MetricCollection
>>> from torchmetrics.wrappers import MultitaskWrapper
>>> from torchmetrics.regression import MeanSquaredError, MeanAbsoluteError
>>> from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
>>>
>>> classification_target = torch.tensor([0, 1, 0])
>>> regression_target = torch.tensor([2.5, 5.0, 4.0])
>>> targets = {"Classification": classification_target, "Regression": regression_target}
>>>
>>> classification_preds = torch.tensor([0, 0, 1])
>>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
>>> preds = {"Classification": classification_preds, "Regression": regression_preds}
>>>
>>> metrics = MultitaskWrapper({
...     "Classification": MetricCollection(BinaryAccuracy(), BinaryF1Score()),
...     "Regression": MetricCollection(MeanSquaredError(), MeanAbsoluteError())
... })
>>> metrics.update(preds, targets)
>>> metrics.compute()
{'Classification': {'BinaryAccuracy': tensor(0.3333), 'BinaryF1Score': tensor(0.)},
 'Regression': {'MeanSquaredError': tensor(0.8333), 'MeanAbsoluteError': tensor(0.6667)}}
```

Example (with a prefix and postfix):

```python
>>> import torch
>>> from torchmetrics.wrappers import MultitaskWrapper
>>> from torchmetrics.regression import MeanSquaredError
>>> from torchmetrics.classification import BinaryAccuracy
>>>
>>> classification_target = torch.tensor([0, 1, 0])
>>> regression_target = torch.tensor([2.5, 5.0, 4.0])
>>> targets = {"Classification": classification_target, "Regression": regression_target}
>>> classification_preds = torch.tensor([0, 0, 1])
>>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
>>> preds = {"Classification": classification_preds, "Regression": regression_preds}
>>>
>>> metrics = MultitaskWrapper({
...     "Classification": BinaryAccuracy(),
...     "Regression": MeanSquaredError()
... }, prefix="train_")
>>> metrics.update(preds, targets)
>>> metrics.compute()
{'train_Classification': tensor(0.3333), 'train_Regression': tensor(0.8333)}
```

#### `clone(prefix=None, postfix=None)` {#torchmetrics.wrappers.MultitaskWrapper.clone}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L276-L287)

Make a copy of the metric.

Parameters:

*   **prefix** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]) – a string to append in front of the metric keys
    
*   **postfix** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]) – a string to append after the keys of the output dict.
    

Return type:

[`MultitaskWrapper`](#torchmetrics.wrappers.MultitaskWrapper "torchmetrics.wrappers.multitask.MultitaskWrapper")

#### `compute()` {#torchmetrics.wrappers.MultitaskWrapper.compute}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L236-L246)

Compute metrics for all tasks.

Return type:

[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]

#### `forward(task_preds, task_targets)` {#torchmetrics.wrappers.MultitaskWrapper.forward}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L248-L257)

Call underlying forward methods for all tasks and return the result as a dictionary.

Return type:

[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]

#### `items(flatten=True)` {#torchmetrics.wrappers.MultitaskWrapper.items}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L163-L176)

Iterate over task and task metrics.

Parameters:

**flatten** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – If True, will iterate over all sub-metrics in the case of a MetricCollection. If False, will iterate over the task names and the corresponding metrics.

Return type:

[`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")\]\]

#### `keys(flatten=True)` {#torchmetrics.wrappers.MultitaskWrapper.keys}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L178-L191)

Iterate over task names.

Parameters:

**flatten** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – If True, will iterate over all sub-metrics in the case of a MetricCollection. If False, will iterate over the task names and the corresponding metrics.

Return type:

[`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]

#### `plot(val=None, axes=None)` {#torchmetrics.wrappers.MultitaskWrapper.plot}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L289-L389)

Plot a single or multiple values from the metric.

All tasks’ results are plotted on individual axes.

Parameters:

*   **val** (`Union`\[[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)"), [`Sequence`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence "(in Python v3.15)")\[[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Either a single result from calling metric.forward or metric.compute or a list of these results. If no value is provided, will automatically call metric.compute and plot that result.
    
*   **axes** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`Sequence`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence "(in Python v3.15)")\[[`Axes`](https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.html#matplotlib.axes.Axes "(in Matplotlib v3.11.2)")\]\]) – Sequence of matplotlib axis objects. If provided, will add the plots to the provided axis objects. If not provided, will create them.
    

Return type:

[`Sequence`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence "(in Python v3.15)")\[[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[[`Figure`](https://matplotlib.org/stable/api/_as_gen/matplotlib.figure.Figure.html#matplotlib.figure.Figure "(in Matplotlib v3.11.2)"), `Union`\[[`Axes`](https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.html#matplotlib.axes.Axes "(in Matplotlib v3.11.2)"), [`ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray "(in NumPy v2.5)")\]\]\]

Returns:

Sequence of tuples with Figure and Axes object for each task.

```python
>>> # Example plotting a single value
>>> import torch
>>> from torchmetrics.wrappers import MultitaskWrapper
>>> from torchmetrics.regression import MeanSquaredError
>>> from torchmetrics.classification import BinaryAccuracy
>>>
>>> classification_target = torch.tensor([0, 1, 0])
>>> regression_target = torch.tensor([2.5, 5.0, 4.0])
>>> targets = {"Classification": classification_target, "Regression": regression_target}
>>>
>>> classification_preds = torch.tensor([0, 0, 1])
>>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
>>> preds = {"Classification": classification_preds, "Regression": regression_preds}
>>>
>>> metrics = MultitaskWrapper({
...     "Classification": BinaryAccuracy(),
...     "Regression": MeanSquaredError()
... })
>>> metrics.update(preds, targets)
>>> value = metrics.compute()
>>> fig_, ax_ = metrics.plot(value)
```

[![../../\_images/multi\_task\_wrapper-1\_00.png](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-1_00.png)](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-1_00.png) [![../../\_images/multi\_task\_wrapper-1\_01.png](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-1_01.png)](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-1_01.png)

```python
>>> # Example plotting multiple values
>>> import torch
>>> from torchmetrics.wrappers import MultitaskWrapper
>>> from torchmetrics.regression import MeanSquaredError
>>> from torchmetrics.classification import BinaryAccuracy
>>>
>>> classification_target = torch.tensor([0, 1, 0])
>>> regression_target = torch.tensor([2.5, 5.0, 4.0])
>>> targets = {"Classification": classification_target, "Regression": regression_target}
>>>
>>> classification_preds = torch.tensor([0, 0, 1])
>>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
>>> preds = {"Classification": classification_preds, "Regression": regression_preds}
>>>
>>> metrics = MultitaskWrapper({
...     "Classification": BinaryAccuracy(),
...     "Regression": MeanSquaredError()
... })
>>> values = []
>>> for _ in range(10):
...     values.append(metrics(preds, targets))
>>> fig_, ax_ = metrics.plot(values)
```

[![../../\_images/multi\_task\_wrapper-2\_00.png](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-2_00.png)](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-2_00.png) [![../../\_images/multi\_task\_wrapper-2\_01.png](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-2_01.png)](/docs/torchmetrics/1.9.0/static/_images/multi_task_wrapper-2_01.png)

#### `reset()` {#torchmetrics.wrappers.MultitaskWrapper.reset}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L259-L268)

Reset all underlying metrics.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `update(task_preds, task_targets)` {#torchmetrics.wrappers.MultitaskWrapper.update}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L207-L230)

Update each task’s metric with its corresponding pred and target.

Parameters:

*   **task\_preds** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – Dictionary associating each task to a Tensor of pred.
    
*   **task\_targets** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – Dictionary associating each task to a Tensor of target.
    

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `values(flatten=True)` {#torchmetrics.wrappers.MultitaskWrapper.values}

[\[source\]](https://github.com/Lightning-AI/torchmetrics/blob/v1.9.0/src/torchmetrics/wrappers/multitask.py#L193-L205)

Iterate over task metrics.

Parameters:

**flatten** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – If True, will iterate over all sub-metrics in the case of a MetricCollection. If False, will iterate over the task names and the corresponding metrics.

Return type:

[`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")\]

---

**Navigation**

- Parent: [Wrappers](/docs/torchmetrics/1.9.0/wrappers.md)