def with_task_labels(obj):
    """Add `TaskAware` attributes.

    The dataset must already have task labels.

    `obj` must be a scenario, stream, or experience.
    """

    def _add_task_labels(exp):
        tls = exp.dataset.targets_task_labels.uniques
        # tls is a set, we need to convert to list to call __getitem__
        tls = list(tls)
        if len(tls) == 1:
            exp.task_label = tls[0]
        exp.task_labels = tls
        return exp

    return _decorate_generic(obj, _add_task_labels)
Add `TaskAware` attributes.

    The dataset must already have task labels.

    `obj` must be a scenario, stream, or experience.
    
avalanche/benchmarks/scenarios/task_aware.py